Bugfix: now works correctly even if the input Tensor(s) do not require grad.

This works by adding a fake input to the forward() of torch.autograd.Function.

Also adds an automatic detection of prologue (first layers with no trainable parameters)