Skip to content
New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Jump to bottom

ViT does not support torchscript #3432

Closed
yiheng-wang-nv opened this issue Dec 2, 2021 · 2 comments · Fixed by #3782
Closed

ViT does not support torchscript #3432

yiheng-wang-nv opened this issue Dec 2, 2021 · 2 comments · Fixed by #3782
Assignees
@yiheng-wang-nv
Labels
enhancement New feature or request

Comments

@yiheng-wang-nv
Copy link
Contributor

yiheng-wang-nv commented Dec 2, 2021 •
edited
Loading

class PatchEmbeddingBlock does not support torchscript, due to the module einops.layers.torch.Rearrange used.

Update (02.09.2022) After Einops v0.4.0 released, it supports torchscript now, but the current ViT implementation still cannot support torchscript directly, some changes are necessary, and the einops usage should be kept.
@yiheng-wang-nv yiheng-wang-nv self-assigned this Dec 2, 2021
@yiheng-wang-nv
Copy link
Contributor Author

see also #3428

@yiheng-wang-nv yiheng-wang-nv added the enhancement New feature or request label Dec 2, 2021
@yiheng-wang-nv yiheng-wang-nv mentioned this issue Dec 2, 2021
Closed
7 tasks
@ahatamiz
Copy link
Contributor

ahatamiz commented Dec 2, 2021

Hi @yiheng-wang-nv

Maybe you look at this:
https://github.com/arogozhnikov/einops/blob/45e795f122840beec70f0e1970b1e37b5809a8e6/tests/test_layers.py#L188

torchscript works on Rearrange.

@yiheng-wang-nv yiheng-wang-nv mentioned this issue Feb 9, 2022
Merged
7 tasks
@yiheng-wang-nv yiheng-wang-nv changed the title PatchEmbeddingBlock does not support torchscript ViT does not support torchscript Feb 9, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@yiheng-wang-nv yiheng-wang-nv

Labels
enhancement New feature or request
Projects
None yet
Milestone
No milestone
2 participants
@ahatamiz @yiheng-wang-nv