@beneyal
After calling model.cuda(). This is my experiment function
def experiment(dict_partition, device, bar, args):
    time_start = time.time()

model = RNNNet(args)
model.cuda()

if args.optim == 'Adam':
    optimizer = optim.Adam(model.parameters(),
                           lr=args.lr,
                           weight_decay=args.l2_coef)
elif args.optim == 'RMSprop':
    optimizer = optim.RMSprop(model.parameters(),
                           lr=args.lr,
                           weight_decay=args.l2_coef)
elif args.optim == 'SGD':
    optimizer = optim.SGD(model.parameters(),
                           lr=args.lr,
                           weight_decay=args.l2_coef)
else:
    assert False, 'Undefined Optimizer Type'

criterion = nn.MSELoss()
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=args.gamma)

list_train_loss = list()
list_val_loss = list()
list_train_mae = list()
list_val_mae = list()

data_train = DataLoader(dict_partition['train'], batch_size=args.batch_size, shuffle=args.shuffle)
data_val = DataLoader(dict_partition['val'], batch_size=args.batch_size, shuffle=args.shuffle)

for epoch in range(args.epoch):
    scheduler.step()
    
    model, train_loss, train_mae = train(model, device, optimizer, criterion, data_train, bar, args)
    list_train_loss.append(train_loss)
    list_train_mae.append(train_mae)

    mode, val_loss, val_mae = validate(model, device, criterion, data_val, bar, args)
    list_val_loss.append(val_loss)
    list_val_mae.append(val_mae)

data_test = DataLoader(dict_partition['test'], batch_size=args.batch_size, shuffle=args.shuffle)

mae, std, logP_total, pred_logP_total = test(model, device, data_test, args)

time_end = time.time()
time_required = time_end - time_start

args.list_train_loss = list_train_loss
args.list_val_loss = list_val_loss
args.list_train_mae = list_train_mae
args.list_val_mae = list_val_mae
args.logP_total = logP_total
args.pred_logP_total = pred_logP_total
args.mae = mae
args.std = std
args.time_required = time_required

return args