From 8ebfdd6ca70f3495f41d181cfaf853918082c6f2 Mon Sep 17 00:00:00 2001 From: maciej Date: Fri, 15 May 2020 15:11:23 +0100 Subject: [PATCH] Fixing a bug of mixing data stored on gpu and a model stored on cpu while using the following command: python main.py --cpu --evaluate [path2checkpoint] --- main.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/main.py b/main.py index 94e6aec..4ec8c6c 100644 --- a/main.py +++ b/main.py @@ -299,7 +299,8 @@ def main(): optimizer.load_state_dict(checkpoint['optimizer']) print("=> checkpoint state loaded.") - model = torch.nn.DataParallel(model) + if cuda: + model = torch.nn.DataParallel(model) # Data loading code print("=> creating data loaders ... ")