Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import torchvision.models as models 
 from torchvision import transforms as T
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 new_model = models.resnet18()
 num_ftrs = new_model.fc.in_features 
 new_model.fc = nn.Linear(num_ftrs, 75)