Make JIT not assume that the device is CUDA. (#54238)
Summary:
Decouple the JIT argument spec and shape analysis with CUDA.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/54238
Reviewed By: ngimel
Differential Revision: D28802085
Pulled By: Krovatkin
fbshipit-source-id: 4068c9460cdec2d80733f001ca90ea3f5e6d3a7e
