Resolved problem where unit tests for inference/train_model didn't work on GPU

qim3d/utils/models.py

@@ -66,13 +66,15 @@ def train_model(model, hyperparameters, train_loader, val_loader, eval_every = 1
                 inputs, targets = data
                 inputs = inputs.to(device)
 
-                if device == 'cuda':
+                if torch.cuda.is_available():
+                    
                     targets = targets.to(device).type(torch.cuda.FloatTensor).unsqueeze(1)
                 else:
                     targets = targets.to(device).type(torch.FloatTensor).unsqueeze(1)
     
                 optimizer.zero_grad()
                 outputs = model(inputs)
+                
                 loss = criterion(outputs, targets)
     
                 # Backpropagation
@@ -99,7 +101,7 @@ def train_model(model, hyperparameters, train_loader, val_loader, eval_every = 1
                     inputs, targets = data
                     inputs = inputs.to(device)
                     
-                    if device == 'cuda':
+                    if torch.cuda.is_available():
                         targets = targets.to(device).type(torch.cuda.FloatTensor).unsqueeze(1)
                     else:
                         targets = targets.to(device).type(torch.FloatTensor).unsqueeze(1)
@@ -207,6 +209,7 @@ def inference(data,model):
     else:
         raise ValueError("Input image must be (C,H,W) format")
 
+    model.to(device)
     model.eval()
 
     # Make new list such that possible augmentations remain identical for all three rows