Former-commit-id: 72f12645410e87ba18093f055c9d8d9f8045595b
This commit is contained in:
milesial 2019-10-30 12:27:03 +01:00 committed by GitHub
parent 4c0f0a7a7b
commit 6f23624412

View file

@ -111,7 +111,7 @@ if __name__ == "__main__":
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}') logging.info(f'Using device {device}')
net.to(deviec=device) net.to(device=device)
net.load_state_dict(torch.load(args.model, map_location=device)) net.load_state_dict(torch.load(args.model, map_location=device))
logging.info("Model loaded !") logging.info("Model loaded !")