sunana commited on
Commit
f1dea0c
1 Parent(s): 6144bbe

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -6
app.py CHANGED
@@ -80,18 +80,18 @@ md = "![](https://drive.google.com/uc?id=1WBqYsKRwn_78A72MJBrk643l3-gfAssP) \n"
80
  "This project is licensed under the terms of the MIT license. \n"
81
 
82
  if __name__ =='__main__':
83
- torch.cuda.init()
84
- print(f"Is CUDA available: {torch.cuda.is_available()}")
85
- # True
86
- print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
87
- # Tesla T4
88
 
89
  model = FFV1MT_MS.FFV1DNN()
90
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
91
 
92
  print('Number fo parameters: {}'.format(model.num_parameters()))
93
  model.to(device)
94
- model_dict = torch.load('Model_example.pth.tar', map_location=device)['state_dict']
95
  # save model
96
  model.load_state_dict(model_dict, strict=True)
97
  model.eval()
 
80
  "This project is licensed under the terms of the MIT license. \n"
81
 
82
  if __name__ =='__main__':
83
+ # torch.cuda.init()
84
+ # print(f"Is CUDA available: {torch.cuda.is_available()}")
85
+ # # True
86
+ # print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
87
+ # # Tesla T4
88
 
89
  model = FFV1MT_MS.FFV1DNN()
90
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
91
 
92
  print('Number fo parameters: {}'.format(model.num_parameters()))
93
  model.to(device)
94
+ model_dict = torch.load('Model_example.pth.tar', map_location="cpu")['state_dict']
95
  # save model
96
  model.load_state_dict(model_dict, strict=True)
97
  model.eval()