Implement Module with Register Buffer for Causal MaskDeep-ML | Problem