Hi,
I think you want torch.nn.functional.log_softmax 66 instead of the module. I’m not sure about the output shape you want (maybe an additional mean or sum somewhere?), but other than that it looks very reasonable to me.
Best regards
Thomas