Hi,
I guess it’s a typo but in the first Function, you forgot to save weight and bias.
a,b and c are not inputs or outputs so you can’t give them to save_for_backward() (it should crash no?). You can save them to ctx.abc = (a,b,c) or ctx.a = a, ctx.b = b etc.
How do you test this code? It is possible that another Function holds onto input and so input is still kept alive.
Do you have a full code sample that reproduce your issue?