Method __init__

train.py:94–102

(self, in_dim=21, h1=64, h2=32)

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
class MLP(nn.Module):
    def __init__(self, in_dim=21, h1=64, h2=32):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, h1),
            nn.ReLU(),
            nn.Linear(h1, h2),
            nn.ReLU(),
            nn.Linear(h2, 1)
        )
    def forward(self, x):
        return self.net(x)