Spaces:
Running
Running
Update prediction.py
Browse files- prediction.py +1 -1
prediction.py
CHANGED
@@ -77,7 +77,7 @@ def load_model():
|
|
77 |
|
78 |
# Load the state_dict (weights) from the saved model file
|
79 |
try:
|
80 |
-
state_dict = torch.load("transformer_model.bin",map_location=device) # Ensure loading on the correct device
|
81 |
model.load_state_dict(state_dict)
|
82 |
model.eval()
|
83 |
model.to(device) # Send model to GPU if available
|
|
|
77 |
|
78 |
# Load the state_dict (weights) from the saved model file
|
79 |
try:
|
80 |
+
state_dict = torch.load("transformer_model(1).bin",map_location=device) # Ensure loading on the correct device
|
81 |
model.load_state_dict(state_dict)
|
82 |
model.eval()
|
83 |
model.to(device) # Send model to GPU if available
|