Skip to content

Commit

Permalink
Merge pull request #3272 from flairNLP/3266-typeerror-state_dict-got-…
Browse files Browse the repository at this point in the history
…multiple-values-for-argument-destination

fix multiple arguments destination
  • Loading branch information
alanakbik authored Aug 8, 2023
2 parents e7a985b + 511ce52 commit 911d915
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 911d915

Please sign in to comment.