• Saving is handled by lightning trainer
• loading too is mainly handled by the lightning module, the method load_from_checkpoint is implemented in torchFastText to load smoothly from a checkpoint path

I'm reopening this, I think there could be some simplification in the load_from_checkpoint method.

At the moment, loading a previously trained model needs the training texts (to build the tokenizer in the state required), the exact config for the training (which can be loaded from json), the checkpoint folder (from lightning), and finally to set the flag trained to True :

torchft_model = torchFastText.from_json('torchft_model')
torchft_model.build(np.asarray(X_train), lr=0.1)
torchft_model.load_from_checkpoint('lightning_logs/version_1/checkpoints/epoch=0-step=1705.ckpt')
torchft_model.trained=True

Did I miss something ?

Omitting the requirement to have the training texts would make the models more portable across environment I think.

Another thing : Loading a model trained with CUDA and using predict gives the following error :

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)

Adding the line self.pytorch_model.to(X.device) in predict() seems to solve this.