this is now enough to fully reproduce the behavior (together with the upper replies)
import torch
import lib.medzoo as medzoo
import argparse
from torch.nn import MSELoss


def get_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--title', default='first attempt at end to end age prediction without prior training')
    parser.add_argument('--batchSz', type=int, default=2)
    parser.add_argument('--dataset_name', type=str, default="fmri_hcp")
    parser.add_argument('--dim', nargs="+", type=int, default=(64, 64, 32))
    parser.add_argument('--nEpochs', type=int, default=1)
    parser.add_argument('--nEpochs_age_prediction', type=int, default=100)
    parser.add_argument('--augmentation', action='store_true', default=False)
    parser.add_argument('--inChannels', type=int, default=1)
    parser.add_argument('--inModalities', type=int, default=1)
    parser.add_argument('--split', default=0.8, type=float, help='Select percentage of training data(default: 0.8)')
    parser.add_argument('--lr', default=1e-2, type=float,
                        help='learning rate (default: 5e-3)')
    parser.add_argument('--lr_age_prediction', default=1e-3, type=float,
                        help='learning rate (default: 5e-3)')
    parser.add_argument('--lrsteps', default=[1,4,12,30,70,120,180,270,350,500,700])
    parser.add_argument('--inpaint_location', default=[40,30,17])
    parser.add_argument('--overfit', default=False)
    parser.add_argument('--norm_method', default='global') # pre_voxel
    parser.add_argument('--terminal_show_freq', default=1000)
    parser.add_argument('--classes', type=int, default=1)
    parser.add_argument('--cuda', action='store_true', default=True)
    parser.add_argument('--context_0', default=False)
    parser.add_argument('--loadData', default=False)
    parser.add_argument('--workers', default=6)
    parser.add_argument('--time_step', default=1)
    parser.add_argument('--subject_num', default=1)
    parser.add_argument('--subject_embedding', default=False)
    parser.add_argument('--predict_age', default=True)
    parser.add_argument('--plot_log_scale', default=True)
    parser.add_argument('--resume', default='', type=str, metavar='PATH',
                        help='path to latest checkpoint (default: none)')
    parser.add_argument('--model', type=str, default='UNET3D_gony',
                        choices=('VNET', 'VNET2', 'UNET3D', 'UNET3D_gony','DENSENET1', 'DENSENET2', 'DENSENET3', 'HYPERDENSENET'))
    parser.add_argument('--opt', type=str, default='adam',
                        choices=('sgd', 'adam', 'rmsprop'))
    parser.add_argument('--log_dir', type=str,
                        default='../runs/')

    args = parser.parse_args()

    return args

args = get_arguments()
input = torch.rand((2,1,64,64,32))
subj = torch.tensor([1,2])
target = torch.tensor([14,5])

model = half_Unet_age()
optimizer = torch.optim.Adam(model.parameters(),lr=0.001)
criterion = MSELoss()
if args.cuda:
    input = input.cuda()
    subj = subj.cuda()
    target = target.cuda()
    target = target.to(dtype=torch.float32)
    model = model.cuda()
model.train()
for batch_idx in range(50):
    optimizer.zero_grad()
    input.requires_grad = True
    output = model(input,subj)
    loss = criterion(output, target)
    loss.backward()
    train_loss = loss.item()
    optimizer.step()
    print('loss: ',loss.detach().item())