script_model fix for TTA

mpelchat04 · Oct 8, 2024 · 708c78d · 708c78d
1 parent dddeafb
commit 708c78d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/utils/script_model.py b/utils/script_model.py
@@ -32,8 +32,8 @@ def forward(self, input):
         shape = input.shape
         B, C = shape[0], shape[1]
         input = (self.max_val - self.min_val) * (input - self.min) / (self.max -self.min) + self.min_val
-        input = (input.view(B, C, -1) - self.mean) / self.std
-        input = input.view(shape)
+        input = (input.reshape(B, C, -1) - self.mean) / self.std
+        input = input.reshape(shape)
         output = self.model_scripted(input.to(self.device))
         if self.from_logits:
             if self.num_classes == 1: