iicd/autodl-projects
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix bugs in .to(cpu)

Browse Source
This commit is contained in:
D-X-Y
2021-03-30 12:25:47 +00:00
parent e5ec43e04a
commit 756218974f
4 changed files with 17 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
lib/trade_models/quant_transformer.py
View File

@@ -143,6 +143,19 @@ class QuantTransformer(Model):
device = "cpu"
self.device = device
self.model.to(self.device)
# move the optimizer
for param in self.train_optimizer.state.values():
# Not sure there are any global tensors in the state dict
if isinstance(param, torch.Tensor):
param.data = param.data.to(device)
if param._grad is not None:
param._grad.data = param._grad.data.to(device)
elif isinstance(param, dict):
for subparam in param.values():
if isinstance(subparam, torch.Tensor):
subparam.data = subparam.data.to(device)
if subparam._grad is not None:
subparam._grad.data = subparam._grad.data.to(device)
def loss_fn(self, pred, label):
mask = ~torch.isnan(label)