Function grad_layer_wrt_loss

sat/sgm/modules/autoencoding/magvit2_pytorch.py:149–150

(loss: Tensor, layer: nn.Parameter)

@autocast(enabled=False)
@beartype
def grad_layer_wrt_loss(loss: Tensor, layer: nn.Parameter):
    return torch_grad(outputs=loss, inputs=layer, grad_outputs=torch.ones_like(loss), retain_graph=True)[0].detach()


# helper decorators