code: def merge_layer_scales(self, layer_scales):
      max_dim = max([s.shape[-1] for s in layer_scales])
      layer_scales = [
          torch.cat(
              (
                  s,
                  torch.zeros(
                      (1, max_dim - s.shape[-1]),
                      device=torch.cuda.current_device(),
                  ),
              ),
              dim=-1,
          )
          if s.shape[-1] < max_dim
          else s
          for s in layer_scales
      ]
      return torch.cat(layer_scales).unsqueeze(0)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/runtime/weight_quantizer.py
imports: import torch
  from ..module_inject.replace_policy import HFBertLayerPolicy, replace_policies
label_desc: Use image node_id 6 for calling the WeightQuantization obj's underlying member method code with example usage: obj.merge_layer_scales(layer_scales) and returns: torch