fix cuda err

model.py

@@ -88,7 +88,7 @@ class EvalNet:
 
     def _create_classifier(self):
         original_T_size = self.ori_T
-        upsample_module = nn.Sequential(
+        return nn.Sequential(
             nn.AdaptiveAvgPool2d((1, None)),  # F -> 1
             nn.ConvTranspose2d(
                 self.out_channel_before_classifier,
@@ -123,8 +123,6 @@ class EvalNet:
             nn.Conv2d(32, self.cls_num, kernel_size=(1, 1)),
         )
 
-        return upsample_module
-
     def _set_channel_outsize(self):  #### get the output size before classifier ####
         conv2d_out_ch = []
         for name, module in self.model.named_modules():
@@ -245,7 +243,7 @@ class t_EvalNet:
     def _create_classifier(self):
         original_T_size = self.ori_T
         self.avgpool = nn.AdaptiveAvgPool2d((1, None))  # F -> 1
-        upsample_module = nn.Sequential(  # nn.AdaptiveAvgPool2d((1, None)), # F -> 1
+        return nn.Sequential(  # nn.AdaptiveAvgPool2d((1, None)), # F -> 1
             nn.ConvTranspose2d(
                 self.hidden_dim, 256, kernel_size=(1, 4), stride=(1, 2), padding=(0, 1)
             ),
@@ -275,8 +273,6 @@ class t_EvalNet:
             nn.Conv2d(32, self.cls_num, kernel_size=(1, 1)),
         )
 
-        return upsample_module
-
     def _set_classifier(self):  #### set custom classifier ####
         if self.type == "vit" or self.type == "swin_transformer":
             self.classifier = self._create_classifier()
@@ -287,10 +283,11 @@ class t_EvalNet:
     def forward(self, x: torch.Tensor):
         if torch.cuda.is_available():
             x = x.cuda()
+            self.class_token = self.class_token.cuda()
 
         if self.type == "vit":
             x = self.model._process_input(x)
-            batch_class_token = self.class_token.expand(x.size(0), -1, -1).cuda()
+            batch_class_token = self.class_token.expand(x.size(0), -1, -1)
             x = torch.cat([batch_class_token, x], dim=1)
             x = self.model.encoder(x)
             x = x[:, 1:].permute(0, 2, 1)