Yes it worked, you saved my day bro! I added a single iteration without training before the definition of optimizer, and it can now properly update now.
def mask_init(train_loader, model):
    print("Initializing mask shape...")
    for i, (input, target) in enumerate(train_loader):
        if args.cpu == False:
            input = input.cuda(async=True)
            target = target.cuda(async=True)
        output = model(input)
        break
    print("Mask shape initiated!")