How to use the torch.optim.Adam function in torch

def train(train_iter, model, args):
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    model.train()

    # print(train_iter)
    for sample in train_iter:

        # print(sample)
        feature, target, score = sample
        if args.cuda:
            feature, target = feature.cuda(), target.cuda()

        optimizer.zero_grad()
        output = model(feature)

        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()

sol_gt = Variable(x['sol_gt'].type(dtype), requires_grad=False)

            predictions = sol(img)
            predictions = transformation_utils.pt_xyrs_2_xyxy(predictions)
            loss = alignment_loss(predictions, sol_gt, x['label_sizes'], alpha_alignment, alpha_backprop)
            sum_loss += loss.data[0]
            steps += 1

        if epoch == 0:
            print "First Validation Step Complete"
            print "Benchmark Validation CER:", sum_loss/steps
            lowest_loss = sum_loss/steps

            sol, lf, hw = init_model(config, sol_dir='current', only_load='sol')

            optimizer = torch.optim.Adam(sol.parameters(), lr=train_config['sol']['learning_rate'])
            optim_path = os.path.join(train_config['snapshot']['current'], "sol_optim.pt")
            if os.path.exists(optim_path):
                print "Loading Optim Settings"
                optimizer.load_state_dict(safe_load.torch_state(optim_path))
            else:
                print "Failed to load Optim Settings"

        elif lowest_loss > sum_loss/steps:
            lowest_loss = sum_loss/steps
            print "Saving Best"

            dirname = train_config['snapshot']['best_validation']
            if not len(dirname) != 0 and os.path.exists(dirname):
                os.makedirs(dirname)

            save_path = os.path.join(dirname, "sol.pt")

def reparameterize(self, mu, logvar):
        normal = ndistributions.Normal(mu, logvar.exp())
        return normal.rsample(), normal

    def decode(self, z):
        return z.op(self.fc3, F.relu).op(self.fc4, x="z").sigmoid()

    def forward(self, x):
        mu, logvar = self.encode(x.stack(x=("ch", "height", "width")))
        z, normal = self.reparameterize(mu, logvar)
        return self.decode(z), normal


model = VAE().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)


# Reconstruction + KL divergence losses summed over all elements and batch
def loss_function(recon_x, x, var_posterior):
    BCE = recon_x.reduce2(
        x.stack(h=("ch", "height", "width")),
        lambda x, y: F.binary_cross_entropy(x, y, reduction="sum"),
        ("batch", "x"),
    )
    prior = ndistributions.Normal(
        ntorch.zeros(dict(batch=1, z=1)), ntorch.ones(dict(batch=1, z=1))
    )
    KLD = ndistributions.kl_divergence(var_posterior, prior).sum()
    return BCE + KLD

if self.optimizer == 'adam':
            options_list = []
            for name, params in self.model.named_parameters():
                if name.endswith('bias'):
                    # Workaround bug: [params] instead of params
                    # https://discuss.pytorch.org/t/problem-on-different-learning-rate-and-weight-decay-in-different-layers/3619
                    options = {'params': [params],
                               'lr': self.step_size,
                               'weight_decay': 0}
                else:  # name.endswith('weight')
                    options = {'params': [params],
                               'lr': self.step_size,
                               'weight_decay': self.alpha}
                options_list.append(options)
            optimizer = torch.optim.Adam(options_list)

            # Train loop
            n_iter = 0
            old_epoch = -1
            epoch = 0

            while epoch < self.max_iter:
                self.model.train()
                optimizer.zero_grad()
                for i, loader_iter in enumerate(loaders_iter):
                    try:
                        X_batch, y_batch = next(loader_iter)
                    except StopIteration:
                        loader_iter = iter(loaders[i])
                        X_batch, y_batch = next(loader_iter)
                        loaders_iter[i] = loader_iter

# Loop over epochs.
lr = args.lr
best_val_loss = []
stored_loss = 100000000

# At any point you can hit Ctrl + C to break out of training early.
try:
    optimizer = None
    # Ensure the optimizer is optimizing params, which includes both the model's weights as well as the criterion's weight (i.e. Adaptive Softmax)
    if args.optimizer == 'sgd':
        optimizer = torch.optim.SGD(params, lr=args.lr, weight_decay=args.wdecay)
    if args.optimizer == 'adagrad':
        optimizer = torch.optim.Adagrad(params, lr=args.lr, weight_decay=args.wdecay)
    if args.optimizer == 'adam':
        optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.wdecay)
    if args.optimizer == 'adamw':
        optimizer = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.wdecay)
    if args.optimizer == 'lamb':
        from pytorch_lamb import Lamb
        optimizer = Lamb(params, lr=args.lr, weight_decay=args.wdecay, min_trust=0.25)
        #optimizer = Lamb(params, lr=args.lr, weight_decay=args.wdecay, min_trust=0.1)
        #optimizer = Lamb(params, lr=args.lr, weight_decay=args.wdecay, min_trust=0, random_min_trust=0.2, random_trust_dice=10)
        #optimizer = Lamb(params, lr=args.lr, weight_decay=args.wdecay, min_trust=0.2, random_min_trust=0.5, random_trust_dice=4)
    from lookahead import Lookahead
    if False:
        k, alpha = 5, 0.8
        print('Lookahead - k {} and alpha {}'.format(k, alpha))
        optimizer = Lookahead(base_optimizer=optimizer, k=k, alpha=alpha)

    from apex import amp
    model, optimizer = amp.initialize(model, optimizer, opt_level='O1')

pose_net.load_state_dict(posenet_weights['state_dict'])
        flow_net.load_state_dict(flownet_weights['state_dict'])
        mask_net.load_state_dict(masknet_weights['state_dict'])


    # import ipdb; ipdb.set_trace()
    cudnn.benchmark = True
    disp_net = torch.nn.DataParallel(disp_net)
    pose_net = torch.nn.DataParallel(pose_net)
    mask_net = torch.nn.DataParallel(mask_net)
    flow_net = torch.nn.DataParallel(flow_net)

    print('=> setting adam solver')

    parameters = chain(disp_net.parameters(), pose_net.parameters(), mask_net.parameters(), flow_net.parameters())
    optimizer = torch.optim.Adam(parameters, args.lr,
                                 betas=(args.momentum, args.beta),
                                 weight_decay=args.weight_decay)

    if args.resume and (args.save_path/'optimizer_checkpoint.pth.tar').exists():
        print("=> loading optimizer from checkpoint")
        optimizer_weights = torch.load(args.save_path/'optimizer_checkpoint.pth.tar')
        optimizer.load_state_dict(optimizer_weights['state_dict'])

    with open(args.save_path/args.log_summary, 'w') as csvfile:
        writer = csv.writer(csvfile, delimiter='\t')
        writer.writerow(['train_loss', 'validation_loss'])

    with open(args.save_path/args.log_full, 'w') as csvfile:
        writer = csv.writer(csvfile, delimiter='\t')
        writer.writerow(['train_loss', 'photo_cam_loss', 'photo_flow_loss', 'explainability_loss', 'smooth_loss'])

def __init__(self, model, model_opt, forward_fn, optimizer, lr_scheduler = None):
		super(ModelRunner, self).__init__()

		self.model   	  = model(model_opt)
		self.optimizer    = optimizer
		self.lr_scheduler = lr_scheduler

		if self.optimizer == 'adam' : 
			self.optimizer = optim.Adam(self.model.parameters() 				, \
										lr=model_opt.lr                         , \
										betas=(model_opt.beta1, model_opt.beta2), \
										weight_decay=model_opt.weight_decay)
		elif self.optimizer == 'sgd':
			self.optimizer = optim.SGD(self.model.parameters() 				, \
										lr=model_opt.lr                     , \
										momentum = model_opt.momentum       , \
										weight_decay=model_opt.weight_decay)
		elif self.optimizer == 'adamax':
			self.optimizer = optim.Adamax(self.model.parameters(), lr = model_opt.lr)

		if lr_scheduler == 'step' : self.lr_scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, model_opt.lr_step_size, model_opt.lr_gamma)

		self.gpu 	 	= model_opt.gpu
		self.forward_fn = forward_fn

g_lr, d_lr = args.g_lr, args.d_lr
    gen_parameters = []
    for key, value in dict(netG.named_parameters()).items():
        if value.requires_grad:
            if 'mapping' in key:
                gen_parameters += [{'params': [value], 'lr': g_lr * 0.1}]
            else:
                gen_parameters += [{'params': [value], 'lr': g_lr}]

    g_optimizer = torch.optim.Adam(gen_parameters, betas=(0, 0.999))

    dis_parameters = []
    for key, value in dict(netD.named_parameters()).items():
        if value.requires_grad:
            dis_parameters += [{'params': [value], 'lr': d_lr}]
    d_optimizer = torch.optim.Adam(dis_parameters, betas=(0, 0.999))

    # make dirs
    if not os.path.exists(args.out_path):
        os.mkdir(args.out_path)
    if not os.path.exists(os.path.join(args.out_path, 'model/')):
        os.mkdir(os.path.join(args.out_path, 'model/'))
    writer = SummaryWriter(os.path.join(args.out_path, 'log'))

    logger = setup_logger("lostGAN", args.out_path, 0)
    logger.info(netG)
    logger.info(netD)

    start_time = time.time()
    vgg_loss = VGGLoss()
    vgg_loss = nn.DataParallel(vgg_loss)
    l1_loss = nn.DataParallel(nn.L1Loss())

def set_optimizer(iekf):
    param_list = [{'params': iekf.initprocesscov_net.parameters(),
                           'lr': lr_initprocesscov_net,
                           'weight_decay': weight_decay_initprocesscov_net}]
    for key, value in lr_mesnet.items():
        param_list.append({'params': getattr(iekf.mes_net, key).parameters(),
                           'lr': value,
                           'weight_decay': weight_decay_mesnet[key]
                           })
    optimizer = torch.optim.Adam(param_list)
    return optimizer