Skip to content

GitLab

"docs/en_US/git@developer.sourcefind.cn:OpenDAS/nni.git" did not exist on "bdb2826e49410de3c1e0c6dffa6192833b34dcaa"
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