Okay so the new setParams functions looks like
def setParams(network,state):
params_dict = dict(network['model'].named_parameters())
params=[]
weights=[]
for key, value in params_dict.items():
if key[-4:] == 'bias':
params += [{'params':value,'weight_decay':0.0}]
else:
params += [{'params': value,'weight_decay':state['weght decay']}]
return params
And I guess that is all the required information that I need to feed in? Learning rate etc are already given to the optimizer constructer as a separate variable. Finally this seems to work but if anyone spots any mistakes in the way I change these parameters please let me know thanks