Method forward

torch/testing/_internal/common_fsdp.py:397–398

(self, x)

        return (torch.rand(4, 8, device=device),)

    def forward(self, x):
        return self.module(x)

    def get_loss(self, input, output):
        loss = output.sum()