TypeError: cannot assign 'torch.cuda.FloatTensor' as parameter 'embedding' (torch.nn.Parameter or None expected)

vision
1

Hi, my model has two parameters, embedding and embedding2. I would like to put embedding2 on gpu and embedding on cpu. Is that possible?

import torch
import torch.multiprocessing as mp


class Reason(torch.nn.Module):
    def __init__(self):
        super(Reason, self).__init__()
        self.embedding = torch.nn.Parameter(torch.zeros(100, 100))
        self.embedding2 = torch.nn.Parameter(torch.zeros(100, 100))

def cuda_mp(model, gpu_id):
    # model.cuda(gpu_id) # this works, but it will put both parameters on gpu
    model.embedding = model.embedding.cuda(gpu_id) # this does not work
    print (gpu_id, model.embedding.device, model.embedding2.device)

if __name__ == '__main__':
    torch.multiprocessing.set_start_method('spawn')

    r = Reason()
    r.share_memory()
    print (r.embedding.device, r.embedding2.device)
    procs = []
    for i in range(3):
        p = mp.Process(target=cuda_mp, args=(r, i))
        procs.append(p)
        p.start()
    for p in procs:
        p.join()