import torch import torch.nn as nn import numpy as np class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.linear = nn.Linear(1, 1) def forward(self, x): return self.linear(x)