Skip to content

GitLab

Unverified Commit 19d6b875 authored by uniartisan's avatar uniartisan Committed by GitHub
Browse files

fix: model is not synced to device when not using CPU (#4652)

parent 51d261e7
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
examples/nas/multi-trial/mnist/search.py
View file @ 19d6b875
...@@ -107,6 +107,7 @@ def evaluate_model(model_cls): ...@@ -107,6 +107,7 @@ def evaluate_model(model_cls):
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)
for epoch in range(3): for epoch in range(3):
# train the model for one epoch # train the model for one epoch
train_epoch(model, device, train_loader, optimizer, epoch) train_epoch(model, device, train_loader, optimizer, epoch)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment