@jjerphan I wouldn't be so sure. I think a PyTorch equivalent, possibly wrapped by skorch would be more efficient, even on CPU, especially with if one would use torch.compile to further remove overhead. I might update this notebook another time when I get the chance to confirm.