Update model.py

model.py

@@ -52,17 +52,16 @@ class MobileVision(nn.Module):
         self.projection2 = nn.Linear(4096, 1536).to(DEVICE).to(DTYPE)
 
     def forward(self, x):
-        with torch.no_grad():
-            resized_img = F.interpolate(x, size=(256, 256), mode='bilinear', align_corners=False)
-            out1 = self.vision(resized_img)
-            x = split_chessboard(x, 2)
-            x = self.vision(x)
-            x = merge_chessboard(x, 2)
-            x = F.interpolate(x, size=(8, 8), mode='area')
-            x = torch.cat([out1, x], dim=1)
-
-            x = x.reshape(x.size(0), x.size(1), -1)
-            x = x.permute(0, 2, 1)
+        resized_img = F.interpolate(x, size=(256, 256), mode='bilinear', align_corners=False)
+        out1 = self.vision(resized_img)
+        x = split_chessboard(x, 2)
+        x = self.vision(x)
+        x = merge_chessboard(x, 2)
+        x = F.interpolate(x, size=(8, 8), mode='area')
+        x = torch.cat([out1, x], dim=1)
+
+        x = x.reshape(x.size(0), x.size(1), -1)
+        x = x.permute(0, 2, 1)
 
         x = self.projection(x)
         x = self.projection2(x)