Commit 5d75f825 authored by Kelly Guo's avatar Kelly Guo Committed by Kelly Guo

Saves checkpoint at first iteration of feature extractor (#363)

# Description

Saves a checkpoint file on the first iteration of the feature extractor
so that a checkpoint is available if inferencing early on in the
training process. Previously, we had to wait until 50k steps for the
checkpoint to be saved.

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- Bug fix (non-breaking change which fixes an issue)


## Screenshots

Please attach before and after screenshots of the change if applicable.

<!--
Example:

| Before | After |
| ------ | ----- |
| _gif/png before_ | _gif/png after_ |

To upload images to a PR -- simply drag and drop an image while in edit
mode and it should upload the image directly. You can then paste that
source into the above before/after sections.
-->

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->
parent a62181b2
......@@ -175,14 +175,14 @@ class FeatureExtractor:
pose_loss.backward()
self.optimizer.step()
self.step_count += 1
if self.step_count % 50000 == 0:
torch.save(
self.feature_extractor.state_dict(),
os.path.join(self.log_dir, f"cnn_{self.step_count}_{pose_loss.detach().cpu().numpy()}.pth"),
)
self.step_count += 1
return pose_loss, predicted_pose
else:
img_input = torch.cat((rgb_img, depth_img, segmentation_img), dim=-1)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment