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()