KFrank:

I’m not sure, but maybe this for loop breaks the graph?


Yes, I think you nailed it down.
Here is a small dummy example:
x = torch.randn(10, requires_grad=True)
y = torch.stack([(x == i).float() for i in range(256)])
print(y.requires_grad)
>  False

@KFrank explained the usage of nn.CrossEntropyLoss really well, so I don’t have anything to add.