Skip to content

GitLab

Commit bf6d0025 authored by VoVAllen's avatar VoVAllen Committed by Minjie Wang
Browse files

[Bugfix] fix gpu manual choice for capsule example (#164)

parent 0ec1a492
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
examples/pytorch/capsule/main.py
View file @ bf6d0025
...@@ -61,7 +61,7 @@ def main(): ...@@ -61,7 +61,7 @@ def main():
torch.manual_seed(args.seed) torch.manual_seed(args.seed)
device = torch.device("cuda:1" if use_cuda else "cpu") device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
train_loader = torch.utils.data.DataLoader( train_loader = torch.utils.data.DataLoader(
......
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