Update Training.py

fixed code to work with data loaders

Training.py

@@ -26,6 +26,8 @@ def train_model(
     num_classes : int = 39,
 
 ):
+
+
     
     # Move model to device
     model.to(device)
@@ -54,18 +56,20 @@ def train_model(
 
 
 
-
+    #----------------------
     # training loop
+    #----------------------
+    
     for epoch in range(n_epochs):
         model.train()
         train_accuracy_fn.reset()
 
         # iterate over all the dataloader's mini-batches
-        for i, (inputs, labels) in enumerate(train_loader):
+        for i, batch in enumerate(train_loader):
 
             # move to GPU memory
-            inputs = inputs.to(device)
-            labels = labels.to(device)
+            inputs = batch["image"].to(device)
+            labels = batch["label"].to(device)
 
             # flatten if not cnn REVISE LATER
             if flatten_input:
@@ -113,9 +117,9 @@ def train_model(
         # The context 'torch.no_grad()' tells pytorch we are not interested in computing 
         # gradients here, so forward pass is more efficient
         with torch.no_grad():
-            for i, (inputs, labels) in enumerate(val_loader):
-                inputs = inputs.to(device)
-                labels = labels.to(device)
+            for i, batch in enumerate(val_loader):
+                inputs = batch["image"].to(device)
+                labels = batch["label"].to(device)
 
                 # flatten if not cnn REVISE LATER
                 if flatten_input:
@@ -130,7 +134,7 @@ def train_model(
         val_accuracies[epoch] = current_accuracy
 
 
-            # keep track of best validation accuracy and save best model so far
+        # keep track of best validation accuracy and save best model so far
         if current_accuracy > best_accuracy:
             best_accuracy = current_accuracy
             torch.save(model.state_dict(), save_path)