| 1845 | |
| 1846 | |
| 1847 | class MlpUsingParameters(nn.Module): |
| 1848 | # MLP that uses layers whose parameters need to be targeted with target_parameters |
| 1849 | def __init__(self, bias=True): |
| 1850 | super().__init__() |
| 1851 | |
| 1852 | self.lin0 = _LinearUsingParameter(10, 20, bias=bias) |
| 1853 | self.relu = nn.ReLU() |
| 1854 | self.drop = nn.Dropout(0.5) |
| 1855 | self.lin1 = _LinearUsingParameter(20, 2, bias=bias) |
| 1856 | self.sm = nn.LogSoftmax(dim=-1) |
| 1857 | self.dtype = torch.float |
| 1858 | |
| 1859 | def forward(self, X): |
| 1860 | X = X.to(self.dtype) |
| 1861 | X = self.lin0(X) |
| 1862 | X = self.relu(X) |
| 1863 | X = self.drop(X) |
| 1864 | X = self.lin1(X) |
| 1865 | X = self.sm(X) |
| 1866 | return X |
| 1867 | |
| 1868 | |
| 1869 | class MockTransformerWrapper: |
no outgoing calls
searching dependent graphs…