I add my_weight  to optimizer like below.
my_weight = nn.Parameter(torch.tensor(args.weight).cuda(), requires_grad=True)
self.optimizer.add_param_group({"params": [my_weight]})

And I want freeze when the iteration is 3 times.
for batch_idx, (data, target) in enumerate(train_loader):
    iteration += 1
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)

    if iteration % 3 == 0:
        !!!!! I want freeze my_weight here !!!!!
    else:
        !!!!! I want unfreeze my_weight here !!!!!

    loss.backward()
    optimizer.step()

How can I do?
Thank you.