Update app.py

app.py

@@ -6,7 +6,7 @@ import torchvision.models as models
 import torch.nn as nn
 
 
-class DummyModel(nn.Module):
+class ColorizingModel(nn.Module):
     def __init__(self):
         super(DummyModel, self).__init__()
 
@@ -81,7 +81,8 @@ class DummyModel(nn.Module):
 
         return dec4
 
-
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = ColorizingModel()
 model_weights = torch.load('model.pth', map_location=device)
 model.load_state_dict(model_weights)
 model = model.to(device)