Solved by mruberry in post #4 
              
              
                I tried running your model on Colab and it appears to be training OK for me. I made a few simplifications: 
import torch
import random
import torch.nn as nn
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

data = load_iris()
X = data['data']
y = data['targ…