Add comment

tests/jax/models/squeezebert/test_squeezebert.py

@@ -46,7 +46,7 @@ def _get_forward_method_kwargs(self) -> Dict[str, jax.Array]:
         params = self._model.init_from_pytorch_statedict(state_dict)
 
         return {
-            "variables": params,
+            "variables": params,  # JAX frameworks have a convention of passing weights as the first argument
             "input_ids": self._get_input_activations(),
             "train": False,
         }