Update model.py

model.py

@@ -8,7 +8,7 @@ def create_effnetb2_model(num_classes:int=3,
     """Creates an EfficientNet-B2 feature extractor model and transforms."""
     weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
     transforms = weights.transforms()
-    model = torchvision.models.efficientnet_b2(weights=weights).to(device)
+    model = torchvision.models.efficientnet_b2(weights=weights)
 
     for param in model.parameters():
         param.requires_grad = False