Hi,
I am posting the main part of the code here.

#new
global full_train_set1
global full_train_set2
global val_set1
global val_set2
full_train_set1=torch.load('path to the featueres1 dict.pt')
full_train_set2=torch.load('path to the featueres2 dict.pt') 

val_set1=torch.load('path to the val features1 dict.pt')
val_set2=torch.load('path to the val feature2 dict.pt')
class MyCustomDataset(Dataset):
    def __init__(self,ids,train):
        self.names = ids
        self.train=train
    def __getitem__(self, index):
        my_id=self.names[int(index)]
        if self.train:
        	feature1=full_train_set1[my_id]
            feature2=full_train_set2[my_id]
        	target=full_train_set1[my_id]['target']
        else:
        	feature1=val_set1[my_id]
            feature2=val_set2[my_id]
        	target=val_set[my_id]['target']
        data=(feature1,feature2)
        return (data,target) 
       
    def __len__(self):
        return len(self.names)



def train(epoch,train_loader,optimizer,criterian):
    print('Training epoch..',epoch)
    model.train()
    torch.cuda.empty_cache()
    train_loss=0
    b_cnt=0
    for data,target in train_loader:
        b_cnt=b_cnt+1
        data=list(data)
        data[0]=data[0].cuda()
        data[1]=data[1].cuda()
        target=target.cuda().float()
        optimizer.zero_grad()
        pred=model(data).cuda()
        loss=criterian(pred,target).cuda()
        train_loss+=math.sqrt(loss.item()) 
        loss.backward()
        optimizer.step()
        torch.cuda.empty_cache()
        gc.collect()
    torch.cuda.empty_cache()
    return train_loss/float(b_cnt)



def run():
    torch.cuda.empty_cache()
    n_epochs=100
    print("Running...")
    global n_batches
    n_batches=100
    global model
    model=my_model().cuda()
    model=nn.DataParallel(model,device_ids=[0,1,2])
    criterian=nn.MSELoss().cuda()
    train_ids=find_dataset(full_train_set1) # train_ids is a list 
    l_rate=0.001
    optimizer=torch.optim.Adam(model.parameters(), lr=l_rate)
    train_loader = torch.utils.data.DataLoader(dataset=MyCustomDataset(train_ids,train=True),
                                                batch_size=n_batches,
                                                shuffle=True)

    for epoch in range(1,n_epochs+1):
        gc.collect()
        train_loss=train(epoch,train_loader,optimizer,criterian)
        print('Epoch :'+str(epoch)+': Train rmse:',train_loss)

run()