load state_dict on device instead of cpu

Author: QasimKhan5x

beginner_source/fgsm_tutorial.py

@@ -198,7 +198,7 @@ def forward(self, x):
 model = Net().to(device)
 
 # Load the pretrained model
-model.load_state_dict(torch.load(pretrained_model, map_location='cpu'))
+model.load_state_dict(torch.load(pretrained_model, map_location=device))
 
 # Set the model in evaluation mode. In this case this is for the Dropout layers
 model.eval()