Commit ·
83f2f41
1
Parent(s): d207639
Allow LSTM loading on CPU
Browse files- inference_lstm.py +1 -1
inference_lstm.py
CHANGED
|
@@ -37,7 +37,7 @@ if __name__ == "__main__":
|
|
| 37 |
# Set device
|
| 38 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 39 |
|
| 40 |
-
model_state = torch.load(args.model_path)
|
| 41 |
|
| 42 |
# Load data first
|
| 43 |
train_data, val_data, test_data = load_data(
|
|
|
|
| 37 |
# Set device
|
| 38 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 39 |
|
| 40 |
+
model_state = torch.load(args.model_path, map_location=device)
|
| 41 |
|
| 42 |
# Load data first
|
| 43 |
train_data, val_data, test_data = load_data(
|