Update test_fastdvdnet.py

state_temp_dict = torch.load(args['model_file']) to state_temp_dict = torch.load(args['model_file'], map_location=device)
m-tassano · Dec 14, 2020 · a842411 · a842411
1 parent 905d28f
commit a842411
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test_fastdvdnet.py b/test_fastdvdnet.py
@@ -78,7 +78,7 @@ def test_fastdvdnet(**args):
 	model_temp = FastDVDnet(num_input_frames=NUM_IN_FR_EXT)
 
 	# Load saved weights
-	state_temp_dict = torch.load(args['model_file'])
+	state_temp_dict = torch.load(args['model_file'], map_location=device)
 	if args['cuda']:
 		device_ids = [0]
 		model_temp = nn.DataParallel(model_temp, device_ids=device_ids).cuda()