Hi Aditya!



 Aditya_Prakash:

However, the program gives parameter().grad = None error which is mostly because the graph is breaking somewhere.


I haven’t looked at your code or error messages in any detail.
But, yes, your “graph is breaking.”


  code = 2*(output > 0.5).long()



code, as a function of output, is not (usefully) differentiable, so,
regardless of what you subsequently do with code, you won’t be
able to back propagate through it.
I have no idea whether this would make sense for your use case,
but one possibility would run as follows:
As I read it, code is calculated to be either 0 or 2.  You could
instead interpret output (processed appropriately, as necessary)
to be the probability that code should be 0 vs. 2, and then use
that probability to form a weighted average of the 0 and 2 entries
in your Q_m2 array.
This will be differentiable and you will be able to backpropagate (but
I’m not saying that it would make sense …).
Best.
K. Frank