feat: train embedding layers only (borisdayma#291)

Author: borisdayma

src/dalle_mini/model/modeling.py

@@ -1395,7 +1395,7 @@ class DalleBart(PretrainedFromWandbMixin, FlaxBartForConditionalGeneration):
     def num_params(self, params=None):
         if params is None:
             params = self.params
-        num_params = jax.tree_map(
+        num_params = jax.tree_util.tree_map(
             lambda param: param.size, flatten_dict(unfreeze(params))
         ).values()
         return sum(list(num_params))