test_dataset.py 文件源码

def module_cls(self):
        """Return a simple module that concatenates its 2 inputs in
        forward step.

        """
        class MyModule(nn.Module):
            def __init__(self):
                super(MyModule, self).__init__()
                self.dense = nn.Linear(20, 2)

            # pylint: disable=arguments-differ
            def forward(self, X0, X1):
                X = torch.cat((X0, X1), 1)
                X = F.softmax(self.dense(X), dim=-1)
                return X

        return MyModule