I think a nicer solution here is to just del/set the attribute so that you don’t get into trouble:
module._parameters[param_key] = memo[p]

# Can become
delattr(module, param_key)
setattr(module, param_key, memo[p])

Taking a step back, we have a new experiemental “stateless” version of Module that might help you here (pytorch/_stateless.py at 8532061bce8da8b5fe9ecce1067ade16793a7ee3 · pytorch/pytorch · GitHub 3)
In particular, you can call your Module with a set of params without having to set them as parameters with functional_call(model, temporary_params, input).