From 6dfe5681ade2631f7a23f167fdb0f1c1dca0ec39 Mon Sep 17 00:00:00 2001
From: Chao Fan <92680886+ChaoFan996@users.noreply.github.com>
Date: Tue, 19 Mar 2024 14:53:22 +0800
Subject: [PATCH] Update deepgaitv2.py

---
 opengait/modeling/models/deepgaitv2.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/opengait/modeling/models/deepgaitv2.py b/opengait/modeling/models/deepgaitv2.py
index 3e66a62..60894e0 100644
--- a/opengait/modeling/models/deepgaitv2.py
+++ b/opengait/modeling/models/deepgaitv2.py
@@ -61,6 +61,7 @@ class DeepGaitV2(BaseModel):
             self.layer3 = SetBlockWrapper(self.layer3)
             self.layer4 = SetBlockWrapper(self.layer4)
 
+        self.mode = mode
         self.FCs = SeparateFCs(16, channels[3], channels[2])
         self.BNNecks = SeparateBNNecks(16, channels[2], class_num=model_cfg['SeparateBNNecks']['class_num'])
 
@@ -93,6 +94,10 @@ class DeepGaitV2(BaseModel):
     def forward(self, inputs):
         ipts, labs, typs, vies, seqL = inputs
 
+        if not self.training and len(labs) != 1 and self.mode != '2d':
+            raise ValueError(
+                'The input size of each GPU must be 1 in testing mode, but got {}!'.format(len(labs)))
+        
         if len(ipts[0].size()) == 4:
             sils = ipts[0].unsqueeze(1)
         else: