hiepnguyen034:

Just out of curiosity, can you clarify a bit on how to push the module to cuda “manually” in the forward method?


Just for the sake of completeness, as I don’t recommend doing it for the aforementioned issues:
def forward(self, inputs):
        x = self.main(inputs)
        x = Flatten()(x)
        lin = nn.Linear(x.size(-1),128).to(device)
        x = lin(x)
        x = nn.Sigmoid()(x)
        return x

This will now create a new linear module in each forward pass, push the parameters to the device, and feed x to it. However, please don’t use this code, except you really want to create new linear layers in each forward pass.