I rerun the training module, and made changes to this line
correct_k = correct[:k].view(-1).float().sum(0)

on getting this error: RuntimeError: view size is not compatible with input tensor’s size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(…) instead.

Solves it by calling .contiguous() before view

,
also in the prediction code I changed this
prediction = int(torch.max(output.data, 1)[0].cpu().numpy())

to this
topk=(1,5)
maxk = max(topk)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
print(pred)

taking the first value of the tensor as my prediction, have tried it on ten images of different classes which I got 8 right predictions. Happy to share this with you, thank you very much Dwight_Foster