Kritajns:

preds = np.squeeze(preds_tensor.numpy()) if not train_on_gpu else np.squeeze(preds_tensor.cpu().numpy())


It’s weird. You have a condition that checks if the train_on_gpu is True or False. So, this error should not happen. Can you check the device attribute of the tensor preds_tensor?