better docstring for the use_pytorch_weights function

Author: init-22

custom_pytorch_jax_converter.py

@@ -7,22 +7,19 @@
 import copy
 from jax.tree_util import tree_map
 
-"""
-Jax default parameter structure:
-dict_keys(['Dense_0', 'Dense_1', 'Dense_2', 'Dense_3', 'Dense_4', 'Dense_5', 'Dense_6', 'Dense_7', 'embedding_table'])
-
-Pytorch stateduct structure:
-dict_keys(['embedding_chunk_0', 'embedding_chunk_1', 'embedding_chunk_2', 'embedding_chunk_3', 'bot_mlp.0.weight', 'bot_mlp.0.bias', 'bot_mlp.2.weight', 'bot_mlp.2.bias', 'bot_mlp.4.weight', 'bot_mlp.4.bias', 'top_mlp.0.weight', 'top_mlp.0.bias', 'top_mlp.2.weight', 'top_mlp.2.bias', 'top_mlp.4.weight', 'top_mlp.4.bias', 'top_mlp.6.weight', 'top_mlp.6.bias', 'top_mlp.8.weight', 'top_mlp.8.bias'])
 
+def use_pytorch_weights(file_name: str):
+    """
+    Jax default parameter structure:
+    dict_keys(['Dense_0', 'Dense_1', 'Dense_2', 'Dense_3', 'Dense_4', 'Dense_5', 'Dense_6', 'Dense_7', 'embedding_table'])
 
+    Pytorch stateduct structure:
+    dict_keys(['embedding_chunk_0', 'embedding_chunk_1', 'embedding_chunk_2', 'embedding_chunk_3', 'bot_mlp.0.weight', 'bot_mlp.0.bias', 'bot_mlp.2.weight', 'bot_mlp.2.bias', 'bot_mlp.4.weight', 'bot_mlp.4.bias', 'top_mlp.0.weight', 'top_mlp.0.bias', 'top_mlp.2.weight', 'top_mlp.2.bias', 'top_mlp.4.weight', 'top_mlp.4.bias', 'top_mlp.6.weight', 'top_mlp.6.bias', 'top_mlp.8.weight', 'top_mlp.8.bias'])
 
-The following function converts the PyTorch weights to the Jax format
-and assigns them to the Jax model parameters.
-The function assumes that the Jax model parameters are already initialized
-and that the PyTorch weights are in the correct format.
-"""
 
-def use_pytorch_weights(file_name: str):
+    The following function converts the PyTorch weights to the Jax format
+    """
+    
     jax_copy = {}
 
     # Load PyTorch state_dict lazily to CPU