Update train_and_save_model.py

train_and_save_model.py

@@ -53,8 +53,14 @@ dataset = CustomDataset(texts, labels)
 data_loader = DataLoader(dataset, batch_size=16, shuffle=True)
 
 # Step 6: Training Loop
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)  # Move model to the correct device
+
 for epoch in range(num_epochs):
     for inputs, targets in data_loader:
+        # Move data to the same device as the model
+        inputs, targets = inputs.to(device), targets.to(device)
+
         # Forward pass
         outputs = model(inputs)
         loss = criterion(outputs, targets)
@@ -63,7 +69,11 @@ for epoch in range(num_epochs):
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
+
+        # Print loss for every batch
+        print(f'Epoch [{epoch+1}/{num_epochs}], Batch Loss: {loss.item():.4f}')
     
+    # Print epoch summary
     print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
 
 # Step 7: Save the Model