Update model.py

model.py

@@ -7,7 +7,7 @@ def create_effnetb2_model(num_classes: int = 3,
                           seed: int = 42):
     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
@@ -17,5 +17,5 @@ def create_effnetb2_model(num_classes: int = 3,
     model.classifier = nn.Sequential(
         nn.Dropout(p=0.3, inplace=True),
         nn.Linear(in_features=1408, out_features=num_classes, bias=True)
-    ).to(device)
+    )
     return model, transforms