def to_bfloat16_train(self):
        r"""
        Cast weight to bf16 and it's trail part for training
        """
        trails = []
        for i in range(len(self.weights)):
            if self.weights[i].dtype == torch.float:
                bf16_w, trail = torch.ops.torch_ipex.split_float_bfloat16(
                    self.weights[i]
                )
            elif self.weights[i].dtype == torch.bfloat16:
                bf16_w = self.weights[i]
                trail = torch.zeros_like(bf16_w, dtype=torch.bfloat16)
            elif self.weights[i].dtype == torch.double:
                bf16_w, trail = torch.ops.torch_ipex.split_float_bfloat16(
                    self.weights[i].float()
                )
            else:
                AssertionError(
                    False
                ), r"MergedEmbeddingBag only support dtypes with bfloat, float and double"
            trails.append(trail)
            self.weights[i] = torch.nn.Parameter(bf16_w)
        self.adagrad_args = self.adagrad_args._replace(bf16_trail=trails)

        Cast weight to bf16 and it's trail part for training
        
intel_extension_for_pytorch/nn/modules/merged_embeddingbag.py