Solved by Alexey_Demyanchuk in post #4 
              
              
                In case you want some (sorry it is lengthy  ) code example: 
import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
transform = trans…