These two classes do the same thing as the cppn_image() function above, but I don’t think that I can use them for what I am trying to do?
class CPPN_Conv(nn.Module):

    def __init__(self, size, num_channels=16, num_layers=6, activ_func=nn.Tanh()):
        super(CPPN_Conv, self).__init__()
        self.size = size
        self.net = self.create_net(num_channels, num_layers, activ_func) 		
				
    def create_net(self, num_channels, num_layers, activ_func):
        net = nn.Sequential()
        net.add_module(str(len(net)), nn.Conv2d(in_channels=2, out_channels=num_channels, kernel_size=1))
        net.add_module(str(len(net)), activ_func)
        for l in range(num_layers - 1):
            net.add_module(str(len(net)), nn.Conv2d(in_channels=num_channels, out_channels=num_channels, kernel_size=1))
            net.add_module(str(len(net)), activ_func)
        net.add_module(str(len(net)), nn.Conv2d(in_channels=num_channels, out_channels=3, kernel_size=1))
        net.add_module(str(len(net)), nn.Sigmoid())
        net.apply(self.cppn_normal)
        return net

    def cppn_normal(self, l):
        if type(l) == nn.Conv2d:        
            nn.init.normal_(l.weight)
			
    def forward(self):
        input = create_input(self.size)
        return self.net(input)
		

cppn_c = CPPN_Conv((512,460))
img = cppn_c()



class CPPN_Linear(nn.Module):

    def __init__(self, size, num_channels=16, num_layers=6, activ_func=nn.Tanh()):
        super(CPPN_Linear, self).__init__()
        if type(size) is not tuple:
            size = (size, size)
        self.input_size = size
        self.net = self.create_net(num_channels, num_layers, activ_func)
                
    def create_net(self, num_channels, num_layers, activ_func):
        net = nn.Sequential()
        net.add_module(str(len(net)), nn.Linear(2, num_channels, bias=True))
        net.add_module(str(len(net)), activ_func)
        for l in range(num_layers - 1):
            net.add_module(str(len(net)), nn.Linear(num_channels, num_channels, bias=False))
            net.add_module(str(len(net)), activ_func)
        net.add_module(str(len(net)), nn.Linear(num_channels, 3, bias=False))
        net.add_module(str(len(net)), nn.Sigmoid())
        net.apply(self.cppn_normal)
        return net
        
    def cppn_normal(self, l):
        if type(l) == nn.Linear:        
            nn.init.normal_(l.weight)

    def create_input(self, input_size): 
        w = torch.arange(0, input_size[1])
        h = torch.arange(0, input_size[0]) 
        w_exp = w.unsqueeze(1).expand((input_size[1], input_size[0])).true_divide(input_size[0]) - 0.5
        h_exp = h.unsqueeze(0).expand((input_size[1], input_size[0])).true_divide(input_size[1]) - 0.5
        return torch.stack((w_exp, h_exp), -1).reshape(input_size[1] * input_size[0], 2) # nn.Linear requires a different input shape

    def forward(self):
        input = self.create_input(self.input_size)
        return self.net(input).reshape(self.input_size[1], self.input_size[0], 3).permute(2,1,0).unsqueeze(0)
        
      
cppn_l = CPPN_Linear((512,460))
img = cppn_l()