Hi,
In your custom Function, since you compute the gradient wrt alpha, the not-None result should be second one no?
To ensure the problem is overflow, I would add some prints to find exactly which op does that in your custom Function. To detect nan in aTensor, you can do t.ne(t).any().