Hey,
First, retain_graph=True is not needed if you use create_graph=True.
Also if you expect that all the parameters were used to compute y, you should not pass allow_unused=True.
Can you share how y is computed based on predict_up, predict_down and predict_cls ?
Because if y only depends on predict_cls, given the structure of your forward, it looks expected that these conv layers wouldn’t get gradients as their result will not be used.