Spaces:

taka-yamakoshi
/

causal-intervention-demo

Sleeping

taka-yamakoshi commited on Mar 22, 2023

Commit

0c71efa

•

1 Parent(s): eb3196d

fix imports

Files changed (1) hide show

custom_modeling_albert_flax.py CHANGED Viewed

@@ -12,7 +12,7 @@ from flax.traverse_util import flatten_dict, unflatten_dict
 from jax import lax
 from transformers import AlbertConfig
-from transformers.models.albert.modeling_flax_albert import FlaxAlbertOnlyMLMHead, FlaxAlbertEmbeddings
 from transformers.modeling_flax_outputs import (
     FlaxBaseModelOutput,
     FlaxBaseModelOutputWithPooling,

 from jax import lax
 from transformers import AlbertConfig
+from transformers.models.albert.modeling_flax_albert import FlaxAlbertOnlyMLMHead, FlaxAlbertEmbeddings, FlaxAlbertPreTrainedModel
 from transformers.modeling_flax_outputs import (
     FlaxBaseModelOutput,
     FlaxBaseModelOutputWithPooling,