Skip to content
New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

allow tensors in several schedulers step() call #8905

Merged
merged 1 commit into from
Jul 20, 2024

Conversation

catwell
Copy link
Contributor

@catwell catwell commented Jul 19, 2024

What does this PR do?

Fix typing of the timesteps argument of schedulers taking an integer to also allow tensors. The timestep is passed to _init_step_index which deals with this case.

I did not change the docstrings to "int or torch.Tensor", I can do it if you want (but most other schedulers that already support tensors do not mention it in the docstring).

Note: I kept this PR minimal and did not try to change any actual code, but the typing of timesteps in schedulers is a bit messy. For instance several schedulers initialize self.timesteps as a floating-point tensor but use an integer tensor in set_timesteps.

Before submitting

Who can review?

Maybe @yiyixuxu since this is about schedulers.

Copy link
Collaborator

@yiyixuxu yiyixuxu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks!

@HuggingFaceDocBuilderDev

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@yiyixuxu yiyixuxu merged commit fe79489 into huggingface:main Jul 20, 2024
13 of 15 checks passed
# for free to join this conversation on GitHub. Already have an account? # to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants