Hello Pytorch community,
I am just getting started out with Pytorch. I am trying to test a trained model. Following the Pytorch tutorial, I have this set up:
1    total = 0
2    correct = 0
3    with torch.no_grad():
4        for data in val_loader:
5            images, labels = data
6            outputs = model(images)
7            _, predicted = torch.max(outputs.data, 1)
8            total += labels.size(0)
9            correct += (predicted == labels).sum().item()

Where val_loader is a DataLoader(), and the model class inherits from the nn.Module.
However, line 6 throws the following error:
$ CUDA_VISIBLE_DEVICES=2,3 python tools/test_net.py
Illegal instruction (core dumped)

I am using Pytorch version 1.1.0, with a CUDA version of 10.2, and cuda driver version 440.33.01. Any help would be much appreciated!
Thank you in advance for your help!