Hi Tassilo!



 TJKlein:

However, the backward function in EMALoss is not called and not gradient is computed.


I’m not sure how you might be calling your custom Function, so
maybe that’s the problem, but, on the surface, your code seems
to be working for me:
>>> torch.__version__
'1.9.0'
>>> _ = torch.manual_seed (2021)
>>> class EMALoss(torch.autograd.Function):
...     @staticmethod
...     def forward(ctx, input, running_ema):
...         ctx.save_for_backward(input, running_ema)
...         input_log_sum_exp = input.exp().mean().log()
...
...         return input_log_sum_exp
...
...     @staticmethod
...     def backward(ctx, grad_output):
...         print ('EMALoss::backward: called')
...         EPS = 1e-6
...         input, running_mean = ctx.saved_tensors
...         grad = grad_output * input.exp().detach() / \
...             (running_mean + EPS) / input.shape[0]
...         return grad, None
...
>>> t = torch.randn (3, requires_grad = True)
>>> l = EMALoss().apply (t, torch.zeros (1))
>>> l
tensor(1.4002, grad_fn=<EMALossBackward>)
>>> l.backward()
EMALoss::backward: called
>>> t.grad
tensor([3282087.7500,  633012.5625,  140842.0469])

Best.
K. Frank