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: