removing unecessary steps in training loop.

qim3d/utils/models.py

@@ -21,6 +21,7 @@ def train_model(model, hyperparameters, train_loader, val_loader, eval_every = 1
         eval_every (int, optional): frequency of model evaluation. Defaults to every epoch.
         print_every (int, optional): frequency of log for model performance. Defaults to every 5 epochs.
         
+
     Returns:
         tuple:
             train_loss (dict): dictionary with average losses and batch losses for training loop.
@@ -65,12 +66,7 @@ def train_model(model, hyperparameters, train_loader, val_loader, eval_every = 1
             for data in train_loader:
                 inputs, targets = data
                 inputs = inputs.to(device)
-
-                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)
+                targets = targets.to(device).unsqueeze(1)
     
                 optimizer.zero_grad()
                 outputs = model(inputs)
@@ -100,11 +96,7 @@ def train_model(model, hyperparameters, train_loader, val_loader, eval_every = 1
                 for data in val_loader:
                     inputs, targets = data
                     inputs = inputs.to(device)
-                    
-                    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)
+                    targets = targets.to(device).unsqueeze(1)
                     
                     with torch.no_grad():
                         outputs = model(inputs)