I used differentiable bilinear sampler 4 from STN 2 during loss calculation in my model as below:
mask = model(input)
output = bilinear_sampler(input, mask)
loss = loss_cal(input, output)

Will the gradient flow automatically when loss.backward() is called?
model.optimizer.zerograd()
loss.backward()
model.optimizer.steps()

Should I write backward() function separately for this model? Thank you.