refactor(pt): Optimize type embedding handling in se_atten

- Replaced direct assignment of `type_embd_data` with a call to `register_buffer` for better memory management.
- Improved clarity by using a temporary variable `embd_tensor` for storing the output of the embedding network before registration.
- Maintained functionality for both one-side and two-side type embeddings, ensuring consistent behavior across modes.

These changes enhance the maintainability and performance of the descriptor model.

Author: OutisLi

deepmd/pt/model/descriptor/se_atten.py

@@ -479,9 +479,8 @@ def type_embedding_compression(self, type_embedding_net: TypeEmbedNet) -> None:
             if self.type_one_side:
                 # One-side: only neighbor types, much simpler!
                 # Precompute for all (ntypes+1) neighbor types
-                self.type_embd_data = self.filter_layers_strip.networks[0](
-                    full_embd
-                ).detach()
+                embd_tensor = self.filter_layers_strip.networks[0](full_embd).detach()
+                self.register_buffer("type_embd_data", embd_tensor)
             else:
                 # Two-side: all (ntypes+1)² type pair combinations
                 # Create [neighbor, center] combinations
@@ -494,9 +493,10 @@ def type_embedding_compression(self, type_embedding_net: TypeEmbedNet) -> None:
                 )
                 # Precompute for all type pairs
                 # Index formula: idx = center_type * nt + neighbor_type
-                self.type_embd_data = self.filter_layers_strip.networks[0](
+                embd_tensor = self.filter_layers_strip.networks[0](
                     two_side_embd
                 ).detach()
+                self.register_buffer("type_embd_data", embd_tensor)
 
     def forward(
         self,