Skip to content

Commit

Permalink
fix multiple arguments destination
Browse files Browse the repository at this point in the history
  • Loading branch information
Benedikt Fuchs committed Aug 7, 2023
1 parent f54e6d6 commit 511ce52
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions flair/embeddings/token.py
Original file line number Diff line number Diff line change
Expand Up @@ -428,14 +428,14 @@ def to_params(self) -> Dict[str, Any]:
"embedding_length": self.__embedding_length,
}

def state_dict(self, *args, destination=None, prefix="", keep_vars=False):
def state_dict(self, *args, **kwargs):
# when loading the old versions from pickle, the embeddings might not be added as pytorch module.
# we do this delayed, when the weights are collected (e.g. for saving), as doing this earlier might
# lead to issues while loading (trying to load weights that weren't stored as python weights and therefore
# not finding them)
if list(self.modules()) == [self]:
self.embedding = self.embedding
return super().state_dict(*args, destination=destination, prefix=prefix, keep_vars=keep_vars)
return super().state_dict(*args, **kwargs)


@register_embeddings
Expand Down

0 comments on commit 511ce52

Please sign in to comment.