--- /dev/null
+#!/usr/bin/env python-for-pytorch
+
+from torch import nn, Tensor
+
+##########
+
+class LazyLinear(nn.Module):
+
+ def __init__(self, out_dim, bias = True):
+ super(LazyLinear, self).__init__()
+ self.out_dim = out_dim
+ self.bias = bias
+ self.core = None
+
+ def forward(self, x):
+ x = x.view(x.size(0), -1)
+
+ if self.core is None:
+ if self.training:
+ self.core = nn.Linear(x.size(1), self.out_dim, self.bias)
+ else:
+ raise RuntimeError('Undefined LazyLinear core in inference mode.')
+
+ return self.core(x)
+
+##########
+
+model = nn.Sequential(nn.Conv2d(1, 8, kernel_size = 5),
+ nn.ReLU(inplace = True),
+ LazyLinear(128),
+ nn.ReLU(inplace = True),
+ nn.Linear(128, 10))
+
+# model.eval()
+
+input = Tensor(100, 1, 32, 32).normal_()
+
+output = model(input)