Only pin memory if using CUDA

ayushkarnawat / profit

Exploring evolutionary protein fitness landscapes

MIT License

1 stars 0 forks source link

Only pin memory if using CUDA #114

Open ayushkarnawat opened 4 years ago

ayushkarnawat commented 4 years ago

Currently, no matter if we use CPU or GPU, we pin memory. Rather, we should only pin memory if performing computations on CUDA cores.

use_cuda = args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
dataloader_kwargs = {'pin_memory': True} if use_cuda else {}