I guess you need to set requires_grad = False in the final layer.
I refereed this - RuntimeError: element 0 of variables does not require grad and does not have a grad_fn when trying to implement Gumbel-ST sampler
Do let me know if it works.
1 Like