jakemdaly:

forward


Found solution:
broke the sequential layers up, then converted the data to and from cuda in the forward pass
class Generator(nn.Module):

    def __init__(self, ):
        super(Generator, self).__init__()
        self.z_dim = 10
        self.x_dim = 784
        self.name = 'mnist/dcgan/g_net'



        self.layer1 = nn.ConvTranspose2d(in_channels=10, out_channels=32, kernel_size=(4,4), stride=2, padding=0).to(device=cuda),
        
        self.activation1 = (nn.BatchNorm2d(32), nn.ReLU())
        
        self.layer2 = nn.ConvTranspose2d(in_channels=32, out_channels=32, kernel_size=(6,6),stride=2, padding=0),
        
        self.activation2 = (nn.BatchNorm2d(32), nn.ReLU())

        self.layer3 = nn.ConvTranspose2d(in_channels=32, out_channels=1, kernel_size=(6,6), stride=2, padding=0),
        
        self.activation3 = nn.Sigmoid()

        

        for m in self.modules():
            if isinstance(m, nn.ConvTranspose2d):
                nn.init.normal_(m.weight)


    def forward(self, x):

        out1_ = self.layer1(x.to(device=cuda))
        out1 = self.activation1(out1_.to(device=cpu))
        out2_ = self.layer2(out1)
        out2 = self.activation2(out2_)
        out3_ = self.layer3(out2)
        out3 = self.activation3(out2_.numpy())

        return(out3)