Method backward

lvdm/common.py:120–142

for x in ctx.input_tensors: if isinstance(x, int): print('-----------------', ctx.run_function)

(ctx, *output_grads)

    @staticmethod
    @torch.cuda.amp.custom_bwd # add this
    def backward(ctx, *output_grads):
        '''
        for x in ctx.input_tensors:
            if isinstance(x, int):
                print('-----------------', ctx.run_function)
        '''
        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
        with torch.enable_grad():
            # Fixes a bug where the first op in run_function modifies the
            # Tensor storage in place, which is not allowed for detach()'d
            # Tensors.
            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
            output_tensors = ctx.run_function(*shallow_copies)
        input_grads = torch.autograd.grad(
            output_tensors,
            ctx.input_tensors + ctx.input_params,
            output_grads,
            allow_unused=True,
        )
        del ctx.input_tensors
        del ctx.input_params
        del output_tensors
        return (None, None) + input_grads