Skip to content

Commit a17ec1e

Browse files
committed
Bugfix #839 - fix save function of CLIPModel
1 parent fa3dbe7 commit a17ec1e

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

sentence_transformers/models/CLIPModel.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,11 @@ def save(self, output_path: str):
7373
with open(os.path.join(output_path, 'config.json'), 'w') as fOut:
7474
json.dump({'model_name': model_name, 'tokenizer_name': tokenizer_name}, fOut, indent=2)
7575

76-
self.model.save(os.path.join(output_path, model_name))
76+
## Save model
77+
torch.save(self.model.state_dict(), os.path.join(output_path, model_name))
78+
79+
##Save tokenizer
80+
self.tokenizer.save(output_path)
7781

7882
@staticmethod
7983
def load(input_path: str):

0 commit comments

Comments
 (0)