Solved by ptrblck in post #2 
              
              
                You should be able to add your custom buffers to the corresponding param_group via: 
# setup
lin = nn.Linear(1, 1)
optimizer = torch.optim.Adam(lin.parameters(), lr=1e-3)

# dummy training
out = lin(torch.randn(1, 1)).backward()
optimizer.step()

# check param_group
print(optimizer.param_groups[0])
…