-model = nn.Sequential(nn.Conv2d(1, 8, kernel_size = 5),
- nn.ReLU(inplace = True),
- LazyLinear(128),
- nn.ReLU(inplace = True),
- nn.Linear(128, 10))
+if __name__ == "__main__":
+ model = nn.Sequential(
+ nn.Conv2d(3, 8, kernel_size=5),
+ nn.ReLU(inplace=True),
+ LazyLinear(128),
+ nn.ReLU(inplace=True),
+ nn.Linear(128, 10),
+ )
+
+ # model.eval()