def METHOD_NAME(self, enable_training=False):
    if self.inference:
        if MegatronLayerPolicy.version == 0:
            METHOD_NAME = self.client_module.METHOD_NAME
        else:
            METHOD_NAME = self.client_module.self_attention
    return METHOD_NAME.query_key_value.weight, \
           METHOD_NAME.query_key_value.bias, \
           METHOD_NAME.dense.weight, \
           METHOD_NAME.dense.bias