katsuchi commited on
Commit
f87e858
·
verified ·
1 Parent(s): e3eafcd

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +1 -1
model.py CHANGED
@@ -150,7 +150,7 @@ class Trainer:
150
 
151
  def predict(self, inputs):
152
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
153
- self.model.load_state_dict(torch.load('best_model.pth', weights_only=True, map_location=torch.device('cpu'))
154
  self.model.to(device)
155
  self.model.eval()
156
  inputs = inputs.to(device)
 
150
 
151
  def predict(self, inputs):
152
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
153
+ self.model.load_state_dict(torch.load('best_model.pth', weights_only=True, map_location=torch.device('cpu')))
154
  self.model.to(device)
155
  self.model.eval()
156
  inputs = inputs.to(device)