We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG]: FreqCacheEmbeddingBag._weight.ProcessGroup is initialized before _weight initialized
### 🐛 Describe the bug

When I init a DLRM with ParallelFreqAwareEmbeddingBag, the bug is reported as the following:
![image](https://user-images.githubusercontent.com/34452939/184055340-c99579bf-c9df-43ed-8552-7c9ae7aadffb.png)


I believe that is because the [ParallelFreqAwareEmbeddingBag](https://github.com/hpcaitech/ColossalAI/blob/039b7ed3bc33173e36c5c4decd41f8d7b1ec0f45/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py#L60) init its `_weight.ProcessGroup` before the `_weight` is initialized.

After I swap above line with its next line, the traceback shows another error:
![image](https://user-images.githubusercontent.com/34452939/184055612-de402e67-ac19-4860-ac12-9622f88a6aca.png)

It looks like some api update issue.

### Environment

_No response_

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py`
Content:
```
1 import torch
2 import torch.nn.functional as F
3 from typing import List, Optional, Iterator, Tuple
4 
5 from .base_embedding import BaseEmbeddingBag
6 from .cache_mgr import CachedParamMgr
7 from torch.nn.parameter import Parameter
8 from .._utils import dual_all_to_all
9 
10 from colossalai.tensor import ColoParameter, ShardSpec, ComputeSpec, ComputePattern, ProcessGroup
11 
12 
13 def get_partition(embedding_dim, rank, world_size) -> Tuple[int, int, bool]:
14     if world_size == 1:
15         return 0, embedding_dim, True
16 
17     assert embedding_dim >= world_size, \
18         f"Embedding dimension {embedding_dim} must be larger than the world size " \
19         f"{world_size} of the process group"
20     chunk_size = embedding_dim // world_size
21     threshold = embedding_dim % world_size
22     # if embedding dim is divisible by world size
23     if threshold == 0:
24         return rank * chunk_size, (rank + 1) * chunk_size, True
25 
26     # align with the split strategy of torch.tensor_split
27     size_list = [chunk_size + 1 if i < threshold else chunk_size for i in range(world_size)]
28     offset = sum(size_list[:rank])
29     return offset, offset + size_list[rank], False
30 
31 
32 class ParallelFreqAwareEmbeddingBag(BaseEmbeddingBag):
33 
34     def __init__(self,
35                  num_embeddings,
36                  embedding_dim,
37                  padding_idx=None,
38                  max_norm=None,
39                  norm_type=2.,
40                  scale_grad_by_freq=False,
41                  sparse=False,
42                  _weight=None,
43                  mode='mean',
44                  include_last_offset=False,
45                  dtype=None,
46                  debug=True):
47         super(ParallelFreqAwareEmbeddingBag,
48               self).__init__(num_embeddings, embedding_dim, padding_idx, max_norm, norm_type, scale_grad_by_freq,
49                              sparse, mode, include_last_offset)
50 
51         self.rank = torch.distributed.get_rank()
52         self.world_size = torch.distributed.get_world_size()
53         self.debug = debug
54 
55         self.partition_start_index, self.partition_end_index, divisible = get_partition(
56             embedding_dim, self.rank, self.world_size)
57         self.embedding_dim_per_partition = self.partition_end_index - self.partition_start_index
58 
59         if _weight is None:
60             self._weight.process_group = ProcessGroup(tp_degree=self.world_size)
61             self._weight = ColoParameter.from_torch_tensor(torch.empty(self.num_embeddings,
62                                                                        self.embedding_dim_per_partition,
63                                                                        device='cpu',
64                                                                        dtype=dtype),
65                                                            requires_grad=True,
66                                                            spec=ShardSpec(dims=[-1], num_partitions=[self.world_size]))
67             self.init_parameters()
68         else:
69             assert isinstance(_weight, ColoParameter), "initialized weight must in type of ColoParameter"
70             self._weight = _weight
71 
72     @property
73     def weight(self):
74         return self.cache_weight_mgr.cpu_weight
75 
76     def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
77         yield 'weight', self.cache_weight_mgr.cuda_cached_weight
78 
79     def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
80         yield self.cache_weight_mgr.cuda_cached_weight
81 
82     @torch.no_grad()
83     def init_parameters(self):
84         self._weight.data.uniform_(-1 / self.num_embeddings, 1 / self.num_embeddings)
85         if self.padding_idx is not None:
86             self._weight[self.padding_idx].fill_(0)
87 
88     def preprocess(self,
89                    cuda_row_num: int,
90                    ids_freq_mapping: Optional[List[int]] = None,
91                    warmup_ratio: float = 0.7,
92                    buffer_size: int = 50_000):
93         self.cache_weight_mgr = CachedParamMgr(self._weight, cuda_row_num, buffer_size=buffer_size)
94         self.cache_weight_mgr.reorder(ids_freq_mapping, warmup_ratio)
95 
96     def forward(self, indices, offsets=None, per_sample_weights=None, shape_hook=None, scatter_dim=0, gather_dim=-1):
97         with torch.no_grad():
98             reorder_ids = self.cache_weight_mgr.prepare_ids(indices)
99 
100         output_shard = F.embedding_bag(reorder_ids, self.cache_weight_mgr.cuda_cached_weight, offsets, self.max_norm,
101                                        self.norm_type, self.scale_grad_by_freq, self.mode, self.sparse,
102                                        per_sample_weights, self.include_last_offset, self.padding_idx)
103 
104         if shape_hook is not None:
105             output_shard = shape_hook(output_shard)
106 
107         output_full = dual_all_to_all(output_shard,
108                                       self._weight.get_process_group(),
109                                       scatter_dim=scatter_dim,
110                                       gather_dim=gather_dim)
111         return output_full
112 
113     @classmethod
114     def from_pretrained(cls,
115                         embedding: torch.Tensor,
116                         freeze: bool = True,
117                         padding_idx: Optional[int] = None,
118                         max_norm: Optional[float] = None,
119                         norm_type: float = 2.,
120                         scale_grad_by_freq: bool = False,
121                         sparse: bool = False,
122                         mode: str = 'mean',
123                         include_last_offset: bool = False,
124                         debug: bool = True,
125                         cuda_row_num: int = 100_000,
126                         ids_freq_mapping: Optional[List[int]] = None,
127                         warmup_ratio: float = 0.7) -> 'ParallelFreqAwareEmbeddingBag':
128         rows, cols = embedding.shape
129         embedding_bag = cls(rows, cols, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse, embedding, mode,
130                             include_last_offset, debug)
131         embedding_bag.preprocess(cuda_row_num, ids_freq_mapping, warmup_ratio)
132         embedding_bag.cache_weight_mgr.cuda_cached_weight.requires_grad_ = not freeze
133         return embedding_bag
134
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py b/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py
--- a/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py
+++ b/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py
@@ -7,7 +7,7 @@
 from torch.nn.parameter import Parameter
 from .._utils import dual_all_to_all
 
-from colossalai.tensor import ColoParameter, ShardSpec, ComputeSpec, ComputePattern, ProcessGroup
+from colossalai.tensor import ColoParameter, ShardSpec, ComputeSpec, ComputePattern, ProcessGroup, ColoTensorSpec
 
 
 def get_partition(embedding_dim, rank, world_size) -> Tuple[int, int, bool]:
@@ -57,13 +57,15 @@
         self.embedding_dim_per_partition = self.partition_end_index - self.partition_start_index
 
         if _weight is None:
-            self._weight.process_group = ProcessGroup(tp_degree=self.world_size)
+            colo_tensor_spec = ColoTensorSpec(pg=ProcessGroup(tp_degree=self.world_size),
+                                              dist_attr=ShardSpec(dims=[-1], num_partitions=[self.world_size]),
+                                              compute_attr=ComputePattern.TP1D)
             self._weight = ColoParameter.from_torch_tensor(torch.empty(self.num_embeddings,
                                                                        self.embedding_dim_per_partition,
                                                                        device='cpu',
                                                                        dtype=dtype),
                                                            requires_grad=True,
-                                                           spec=ShardSpec(dims=[-1], num_partitions=[self.world_size]))
+                                                           spec=colo_tensor_spec)
             self.init_parameters()
         else:
             assert isinstance(_weight, ColoParameter), "initialized weight must in type of ColoParameter"

{"golden_diff": "diff --git a/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py b/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py\n--- a/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py\n+++ b/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py\n@@ -7,7 +7,7 @@\n from torch.nn.parameter import Parameter\n from .._utils import dual_all_to_all\n \n-from colossalai.tensor import ColoParameter, ShardSpec, ComputeSpec, ComputePattern, ProcessGroup\n+from colossalai.tensor import ColoParameter, ShardSpec, ComputeSpec, ComputePattern, ProcessGroup, ColoTensorSpec\n \n \n def get_partition(embedding_dim, rank, world_size) -> Tuple[int, int, bool]:\n@@ -57,13 +57,15 @@\n         self.embedding_dim_per_partition = self.partition_end_index - self.partition_start_index\n \n         if _weight is None:\n-            self._weight.process_group = ProcessGroup(tp_degree=self.world_size)\n+            colo_tensor_spec = ColoTensorSpec(pg=ProcessGroup(tp_degree=self.world_size),\n+                                              dist_attr=ShardSpec(dims=[-1], num_partitions=[self.world_size]),\n+                                              compute_attr=ComputePattern.TP1D)\n             self._weight = ColoParameter.from_torch_tensor(torch.empty(self.num_embeddings,\n                                                                        self.embedding_dim_per_partition,\n                                                                        device='cpu',\n                                                                        dtype=dtype),\n                                                            requires_grad=True,\n-                                                           spec=ShardSpec(dims=[-1], num_partitions=[self.world_size]))\n+                                                           spec=colo_tensor_spec)\n             self.init_parameters()\n         else:\n             assert isinstance(_weight, ColoParameter), \"initialized weight must in type of ColoParameter\"\n", "issue": "[BUG]: FreqCacheEmbeddingBag._weight.ProcessGroup is initialized before _weight initialized\n### \ud83d\udc1b Describe the bug\n\nWhen I init a DLRM with ParallelFreqAwareEmbeddingBag, the bug is reported as the following:\r\n![image](https://user-images.githubusercontent.com/34452939/184055340-c99579bf-c9df-43ed-8552-7c9ae7aadffb.png)\r\n\r\n\r\nI believe that is because the [ParallelFreqAwareEmbeddingBag](https://github.com/hpcaitech/ColossalAI/blob/039b7ed3bc33173e36c5c4decd41f8d7b1ec0f45/colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py#L60) init its `_weight.ProcessGroup` before the `_weight` is initialized.\r\n\r\nAfter I swap above line with its next line, the traceback shows another error:\r\n![image](https://user-images.githubusercontent.com/34452939/184055612-de402e67-ac19-4860-ac12-9622f88a6aca.png)\r\n\r\nIt looks like some api update issue.\n\n### Environment\n\n_No response_\n", "before_files": [{"content": "import torch\nimport torch.nn.functional as F\nfrom typing import List, Optional, Iterator, Tuple\n\nfrom .base_embedding import BaseEmbeddingBag\nfrom .cache_mgr import CachedParamMgr\nfrom torch.nn.parameter import Parameter\nfrom .._utils import dual_all_to_all\n\nfrom colossalai.tensor import ColoParameter, ShardSpec, ComputeSpec, ComputePattern, ProcessGroup\n\n\ndef get_partition(embedding_dim, rank, world_size) -> Tuple[int, int, bool]:\n    if world_size == 1:\n        return 0, embedding_dim, True\n\n    assert embedding_dim >= world_size, \\\n        f\"Embedding dimension {embedding_dim} must be larger than the world size \" \\\n        f\"{world_size} of the process group\"\n    chunk_size = embedding_dim // world_size\n    threshold = embedding_dim % world_size\n    # if embedding dim is divisible by world size\n    if threshold == 0:\n        return rank * chunk_size, (rank + 1) * chunk_size, True\n\n    # align with the split strategy of torch.tensor_split\n    size_list = [chunk_size + 1 if i < threshold else chunk_size for i in range(world_size)]\n    offset = sum(size_list[:rank])\n    return offset, offset + size_list[rank], False\n\n\nclass ParallelFreqAwareEmbeddingBag(BaseEmbeddingBag):\n\n    def __init__(self,\n                 num_embeddings,\n                 embedding_dim,\n                 padding_idx=None,\n                 max_norm=None,\n                 norm_type=2.,\n                 scale_grad_by_freq=False,\n                 sparse=False,\n                 _weight=None,\n                 mode='mean',\n                 include_last_offset=False,\n                 dtype=None,\n                 debug=True):\n        super(ParallelFreqAwareEmbeddingBag,\n              self).__init__(num_embeddings, embedding_dim, padding_idx, max_norm, norm_type, scale_grad_by_freq,\n                             sparse, mode, include_last_offset)\n\n        self.rank = torch.distributed.get_rank()\n        self.world_size = torch.distributed.get_world_size()\n        self.debug = debug\n\n        self.partition_start_index, self.partition_end_index, divisible = get_partition(\n            embedding_dim, self.rank, self.world_size)\n        self.embedding_dim_per_partition = self.partition_end_index - self.partition_start_index\n\n        if _weight is None:\n            self._weight.process_group = ProcessGroup(tp_degree=self.world_size)\n            self._weight = ColoParameter.from_torch_tensor(torch.empty(self.num_embeddings,\n                                                                       self.embedding_dim_per_partition,\n                                                                       device='cpu',\n                                                                       dtype=dtype),\n                                                           requires_grad=True,\n                                                           spec=ShardSpec(dims=[-1], num_partitions=[self.world_size]))\n            self.init_parameters()\n        else:\n            assert isinstance(_weight, ColoParameter), \"initialized weight must in type of ColoParameter\"\n            self._weight = _weight\n\n    @property\n    def weight(self):\n        return self.cache_weight_mgr.cpu_weight\n\n    def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:\n        yield 'weight', self.cache_weight_mgr.cuda_cached_weight\n\n    def parameters(self, recurse: bool = True) -> Iterator[Parameter]:\n        yield self.cache_weight_mgr.cuda_cached_weight\n\n    @torch.no_grad()\n    def init_parameters(self):\n        self._weight.data.uniform_(-1 / self.num_embeddings, 1 / self.num_embeddings)\n        if self.padding_idx is not None:\n            self._weight[self.padding_idx].fill_(0)\n\n    def preprocess(self,\n                   cuda_row_num: int,\n                   ids_freq_mapping: Optional[List[int]] = None,\n                   warmup_ratio: float = 0.7,\n                   buffer_size: int = 50_000):\n        self.cache_weight_mgr = CachedParamMgr(self._weight, cuda_row_num, buffer_size=buffer_size)\n        self.cache_weight_mgr.reorder(ids_freq_mapping, warmup_ratio)\n\n    def forward(self, indices, offsets=None, per_sample_weights=None, shape_hook=None, scatter_dim=0, gather_dim=-1):\n        with torch.no_grad():\n            reorder_ids = self.cache_weight_mgr.prepare_ids(indices)\n\n        output_shard = F.embedding_bag(reorder_ids, self.cache_weight_mgr.cuda_cached_weight, offsets, self.max_norm,\n                                       self.norm_type, self.scale_grad_by_freq, self.mode, self.sparse,\n                                       per_sample_weights, self.include_last_offset, self.padding_idx)\n\n        if shape_hook is not None:\n            output_shard = shape_hook(output_shard)\n\n        output_full = dual_all_to_all(output_shard,\n                                      self._weight.get_process_group(),\n                                      scatter_dim=scatter_dim,\n                                      gather_dim=gather_dim)\n        return output_full\n\n    @classmethod\n    def from_pretrained(cls,\n                        embedding: torch.Tensor,\n                        freeze: bool = True,\n                        padding_idx: Optional[int] = None,\n                        max_norm: Optional[float] = None,\n                        norm_type: float = 2.,\n                        scale_grad_by_freq: bool = False,\n                        sparse: bool = False,\n                        mode: str = 'mean',\n                        include_last_offset: bool = False,\n                        debug: bool = True,\n                        cuda_row_num: int = 100_000,\n                        ids_freq_mapping: Optional[List[int]] = None,\n                        warmup_ratio: float = 0.7) -> 'ParallelFreqAwareEmbeddingBag':\n        rows, cols = embedding.shape\n        embedding_bag = cls(rows, cols, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse, embedding, mode,\n                            include_last_offset, debug)\n        embedding_bag.preprocess(cuda_row_num, ids_freq_mapping, warmup_ratio)\n        embedding_bag.cache_weight_mgr.cuda_cached_weight.requires_grad_ = not freeze\n        return embedding_bag\n", "path": "colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py"}], "after_files": [{"content": "import torch\nimport torch.nn.functional as F\nfrom typing import List, Optional, Iterator, Tuple\n\nfrom .base_embedding import BaseEmbeddingBag\nfrom .cache_mgr import CachedParamMgr\nfrom torch.nn.parameter import Parameter\nfrom .._utils import dual_all_to_all\n\nfrom colossalai.tensor import ColoParameter, ShardSpec, ComputeSpec, ComputePattern, ProcessGroup, ColoTensorSpec\n\n\ndef get_partition(embedding_dim, rank, world_size) -> Tuple[int, int, bool]:\n    if world_size == 1:\n        return 0, embedding_dim, True\n\n    assert embedding_dim >= world_size, \\\n        f\"Embedding dimension {embedding_dim} must be larger than the world size \" \\\n        f\"{world_size} of the process group\"\n    chunk_size = embedding_dim // world_size\n    threshold = embedding_dim % world_size\n    # if embedding dim is divisible by world size\n    if threshold == 0:\n        return rank * chunk_size, (rank + 1) * chunk_size, True\n\n    # align with the split strategy of torch.tensor_split\n    size_list = [chunk_size + 1 if i < threshold else chunk_size for i in range(world_size)]\n    offset = sum(size_list[:rank])\n    return offset, offset + size_list[rank], False\n\n\nclass ParallelFreqAwareEmbeddingBag(BaseEmbeddingBag):\n\n    def __init__(self,\n                 num_embeddings,\n                 embedding_dim,\n                 padding_idx=None,\n                 max_norm=None,\n                 norm_type=2.,\n                 scale_grad_by_freq=False,\n                 sparse=False,\n                 _weight=None,\n                 mode='mean',\n                 include_last_offset=False,\n                 dtype=None,\n                 debug=True):\n        super(ParallelFreqAwareEmbeddingBag,\n              self).__init__(num_embeddings, embedding_dim, padding_idx, max_norm, norm_type, scale_grad_by_freq,\n                             sparse, mode, include_last_offset)\n\n        self.rank = torch.distributed.get_rank()\n        self.world_size = torch.distributed.get_world_size()\n        self.debug = debug\n\n        self.partition_start_index, self.partition_end_index, divisible = get_partition(\n            embedding_dim, self.rank, self.world_size)\n        self.embedding_dim_per_partition = self.partition_end_index - self.partition_start_index\n\n        if _weight is None:\n            colo_tensor_spec = ColoTensorSpec(pg=ProcessGroup(tp_degree=self.world_size),\n                                              dist_attr=ShardSpec(dims=[-1], num_partitions=[self.world_size]),\n                                              compute_attr=ComputePattern.TP1D)\n            self._weight = ColoParameter.from_torch_tensor(torch.empty(self.num_embeddings,\n                                                                       self.embedding_dim_per_partition,\n                                                                       device='cpu',\n                                                                       dtype=dtype),\n                                                           requires_grad=True,\n                                                           spec=colo_tensor_spec)\n            self.init_parameters()\n        else:\n            assert isinstance(_weight, ColoParameter), \"initialized weight must in type of ColoParameter\"\n            self._weight = _weight\n\n    @property\n    def weight(self):\n        return self.cache_weight_mgr.cpu_weight\n\n    def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:\n        yield 'weight', self.cache_weight_mgr.cuda_cached_weight\n\n    def parameters(self, recurse: bool = True) -> Iterator[Parameter]:\n        yield self.cache_weight_mgr.cuda_cached_weight\n\n    @torch.no_grad()\n    def init_parameters(self):\n        self._weight.data.uniform_(-1 / self.num_embeddings, 1 / self.num_embeddings)\n        if self.padding_idx is not None:\n            self._weight[self.padding_idx].fill_(0)\n\n    def preprocess(self,\n                   cuda_row_num: int,\n                   ids_freq_mapping: Optional[List[int]] = None,\n                   warmup_ratio: float = 0.7,\n                   buffer_size: int = 50_000):\n        self.cache_weight_mgr = CachedParamMgr(self._weight, cuda_row_num, buffer_size=buffer_size)\n        self.cache_weight_mgr.reorder(ids_freq_mapping, warmup_ratio)\n\n    def forward(self, indices, offsets=None, per_sample_weights=None, shape_hook=None, scatter_dim=0, gather_dim=-1):\n        with torch.no_grad():\n            reorder_ids = self.cache_weight_mgr.prepare_ids(indices)\n\n        output_shard = F.embedding_bag(reorder_ids, self.cache_weight_mgr.cuda_cached_weight, offsets, self.max_norm,\n                                       self.norm_type, self.scale_grad_by_freq, self.mode, self.sparse,\n                                       per_sample_weights, self.include_last_offset, self.padding_idx)\n\n        if shape_hook is not None:\n            output_shard = shape_hook(output_shard)\n\n        output_full = dual_all_to_all(output_shard,\n                                      self._weight.get_process_group(),\n                                      scatter_dim=scatter_dim,\n                                      gather_dim=gather_dim)\n        return output_full\n\n    @classmethod\n    def from_pretrained(cls,\n                        embedding: torch.Tensor,\n                        freeze: bool = True,\n                        padding_idx: Optional[int] = None,\n                        max_norm: Optional[float] = None,\n                        norm_type: float = 2.,\n                        scale_grad_by_freq: bool = False,\n                        sparse: bool = False,\n                        mode: str = 'mean',\n                        include_last_offset: bool = False,\n                        debug: bool = True,\n                        cuda_row_num: int = 100_000,\n                        ids_freq_mapping: Optional[List[int]] = None,\n                        warmup_ratio: float = 0.7) -> 'ParallelFreqAwareEmbeddingBag':\n        rows, cols = embedding.shape\n        embedding_bag = cls(rows, cols, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse, embedding, mode,\n                            include_last_offset, debug)\n        embedding_bag.preprocess(cuda_row_num, ids_freq_mapping, warmup_ratio)\n        embedding_bag.cache_weight_mgr.cuda_cached_weight.requires_grad_ = not freeze\n        return embedding_bag\n", "path": "colossalai/nn/_ops/cache_embedding/parallel_freq_aware_embedding.py"}]}