OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity
ced90d23d3
pytorch/test/cpp/aoti_inference
History
Julius Herb 8f54e56e62 Add optional device index to AOTIModelPackageLoader (#152093) 
This is my suggestion for resolving #152087

This PR extends the constructor of `AOTIModelPackageLoader` with an (optional) device index. The device type is still determined by `metadata_["AOTI_DEVICE_KEY"]`, but the `device_index` argument can be used to move an AOTI model package to different devices like `cuda:0`, `cuda:1`, ... in a convenient way. AFAIK, this is not possible so far using `AOTIModelPackageLoader` alone. The default case (no device index specified) with `metadata_["AOTI_DEVICE_KEY"] == "cuda"` would lead to the current behavior, i.e., the model is loaded to device `cuda`.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/152093
Approved by: https://github.com/desertfire
2025-05-04 11:40:12 +00:00
..
aoti_custom_class.cpp [AOTI] Fix #140546 and support AOTI package load for Intel GPU. (#140664) 2024-12-10 05:05:08 +00:00
aoti_custom_class.h
CMakeLists.txt [ROCm][Inductor] Enable AOT Inductor CPP UTs for ROCm (#131521) 2024-08-08 19:49:56 +00:00
compile_model.py [AOTI] Fix test_aoti_inference CPU build issue (#134675) 2024-08-28 17:42:19 +00:00
generate_lowered_cpu.py [AOTInductor] Add standalone test for compilation from ExportedProgram (#142327) 2024-12-10 06:50:09 +00:00
standalone_compile.sh [AOTInductor] Add standalone test for compilation from ExportedProgram (#142327) 2024-12-10 06:50:09 +00:00
standalone_test.cpp [AOTInductor] Add standalone test for compilation from ExportedProgram (#142327) 2024-12-10 06:50:09 +00:00
test.cpp Add optional device index to AOTIModelPackageLoader (#152093) 2025-05-04 11:40:12 +00:00
test.py [AOTInductor] Free folded constants that's managed by AOTInductor (#149825) 2025-03-27 06:05:50 +00:00