updated app to support CPU

utils.py

@@ -12,12 +12,16 @@ def save_checkpoint(model, optimizer, epoch, loss, path):
     print(f"Checkpoint saved at epoch {epoch}")
 
 def load_checkpoint(model, optimizer, checkpoint_path):
-    checkpoint = torch.load(checkpoint_path, weights_only=True)
+    # Use map_location to load the checkpoint on CPU if CUDA is not available
+    map_location = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    
     model.load_state_dict(checkpoint['model_state_dict'])
     if optimizer is not None:
         optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
     start_epoch = checkpoint['epoch']
     loss = checkpoint['loss']
+    
     return model, optimizer, start_epoch, loss
 
 def plot_training_curves(epochs, train_acc1, test_acc1, train_acc5, test_acc5, train_losses, test_losses, learning_rates):