3 from torch import nn, Tensor
5 ######################################################################
7 class LazyLinear(nn.Module):
9 def __init__(self, out_dim, bias = True):
10 super(LazyLinear, self).__init__()
11 self.out_dim = out_dim
16 x = x.view(x.size(0), -1)
20 self.core = nn.Linear(x.size(1), self.out_dim, self.bias)
22 raise RuntimeError('Undefined LazyLinear core in inference mode.')
26 ######################################################################
28 model = nn.Sequential(nn.Conv2d(1, 8, kernel_size = 5),
29 nn.ReLU(inplace = True),
31 nn.ReLU(inplace = True),
36 input = Tensor(100, 1, 32, 32).normal_()