facebookresearch / facebookresearch/flow_matching

Resuming training from a checkpoint: Incorrect function call and misspelt attribute in ckpt loading logic

Open
#71 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
4.7k
Forks
373
PR merge metrics
No merged PRs in 30d

Description

**Describe the bug**
The checkpoint loading logic contains two bugs due to incorrect function and attribute references, preventing proper resumption of training from a saved checkpoint.

1. In `examples/text/logic/state.py`, [line 62](https://github.com/facebookresearch/flow_matching/blob/c056dd6e4242e9bfb7e43ff615d4da248aaa5e72/examples/text/logic/state.py#L62C13-L62C80): `self._data_state.test.load_state_dict(loaded_state["test_sampler"])`
(FIX: `self._data_state.test.sampler.load_state_dict(loaded_state["test_sampler"])` )
Here, there's a typo that tries to yoink a `state_dict` from a `Dataset` class

2. In `examples/text/main_train.py`, [line 27](https://github.com/facebookresearch/flow_matching/blob/c056dd6e4242e9bfb7e43ff615d4da248aaa5e72/examples/text/run_train.py#L27): `cfg = checkpointing.load_hydra_config_from_run(cfg.load_dir)`
(FIX: cfg = `checkpointing.load_cfg_from_path(cfg.load_dir)` )
Here, the function name is incorrect; the function exists with a different name in `utils/checkpointing.py`.

**To Reproduce**
Set load_dir = 'path_to_ckpt_parent' in `examples/text/configs/config.yaml` and run `examples/text/run_train.py`

**Expected behavior**
The checkpoint gets picked up, and training resumes.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.