This behavior is expected. The reason for this is in the way lambda functions work. The function is defined in the line you marked as not working, but not called. Your variable “name” is changed after the definition and before the call. The last known definition for the variable will then be used at call. To avoid this behavior, you can simply give the lambda function a second variable with a default value, which is set during definition instead of at call. It looks like so:
module.weight.register_hook(lambda x, name=name: x*mask_dict[name])

Then, your code should work just fine.