| 24 | class TestJointGraph(unittest.TestCase): |
| 25 | def test_joint_graph(self) -> None: |
| 26 | class Module(torch.nn.Module): |
| 27 | def __init__(self): |
| 28 | super().__init__() |
| 29 | self.linear = torch.nn.Linear(3, 3) |
| 30 | self.linear_no_train = torch.nn.Linear(3, 3) |
| 31 | for param in self.linear_no_train.parameters(): |
| 32 | param.requires_grad = False |
| 33 | self.loss = torch.nn.CrossEntropyLoss() |
| 34 | |
| 35 | def forward(self, x, y): |
| 36 | return self.loss(self.linear_no_train(self.linear(x)).softmax(dim=0), y) |
| 37 | |
| 38 | m = Module() |
| 39 | example_inputs = (torch.ones(3), torch.tensor([1.0, 0.0, 0.0])) |
no outgoing calls