A problem about using cuda()

[qzhiyue]

Issue Description

os.environ['TL_BACKEND'] = 'torch'
and using gpu0
When I try to put my model and variable into cuda:

net.cuda()
x, y, y_cls = Variable(x).cuda(), Variable(y).cuda(), Variable(y_cls).cuda()

an error occurred here:

class PSPUpsample(tlx.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = tlx.nn.Sequential(
            tlx.nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1, data_format='channels_first'),
            tlx.nn.BatchNorm2d(num_features=out_channels, data_format='channels_first'),
            tlx.layers.PRelu()
        )
    def forward(self, x):
        h, w = 2 * x.size(2), 2 * x.size(3)
        p = tlx.layers.UpSampling2d(scale=(h / x.size(2), w / x.size(3)), method='bilinear',
                                    data_format='channels_first')(x)
        return self.conv(p) #here

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument weight in method wrapper__prelu)

When I put the 'p' into 'cuda:0' before 'return', like this:

p = p.to_device('cuda:0')

the error also occurred.
When I put the 'p' into cpu, it shows that the model in gpu and it doesn't work.

[hanjr92]

[qzhiyue]

This solution is great, thank you!