Pretty sure the issue here is you’re passing in a tensor for the stats on normalize, try passing in just two arrays:
Normalize.from_stats([0.5,0.5,0.5], [0.5,0.5,0.5])
What’s happening here is fastai runs a very basic division when Normalizing and assumes just raw numbers are being passed in, while your tensor here starts out on the CPU. Hence why we have this GPU/CPU cross. You could also just ensure those torch.tensors
are on cuda
(IE torch.tensor([0.5...]).cuda())
)