gistfile1.txt

# (c) Facebook, Inc. and its affiliates. Confidential and proprietary.

import torch
import torch.nn as nn
from torch.autograd import Function


class PassThrough(Function):
    @staticmethod
    def forward(ctx, *inputs):
        return inputs

    @staticmethod
    def backward(ctx, *grad_outputs):
        print(f"grad_outputs {grad_outputs}")
        return grad_outputs


class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Linear(1, 1, bias=False)
        self.b = nn.Linear(1, 1, bias=False)

    def forward(self, x):
        a, b = self.a(x), self.b(x)
        ret = a, b
        new_a, new_b = PassThrough.apply(a, b)
        ret = new_a, new_b
        return ret

model = MyModel()

def print_grads():
    for param_name, param in model.named_parameters():
        print(f"{param_name} : {param.grad}")


inp = torch.ones(1)
print("-- before backward ---")
model.zero_grad()
print_grads()

for _ in range(3):
    model.zero_grad()
    out = model(inp)
    loss = out[0].sum()
    print("Calling backward...")
    loss.backward()
    print("-- after bwd --")
    print_grads()