amitke commited on
Commit
d56edb9
·
1 Parent(s): c3e1f81
Files changed (1) hide show
  1. train.py +1 -0
train.py CHANGED
@@ -14,6 +14,7 @@ import torch.optim as optim
14
 
15
  from models import StockLSTM
16
 
 
17
  ARTIFACTS_DIR = "artifacts"
18
  os.makedirs(ARTIFACTS_DIR, exist_ok=True)
19
 
 
14
 
15
  from models import StockLSTM
16
 
17
+ os.environ["CUDA_VISIBLE_DEVICES"] = ""
18
  ARTIFACTS_DIR = "artifacts"
19
  os.makedirs(ARTIFACTS_DIR, exist_ok=True)
20