frozenleaves/DeepSpeed
1
0
Fork 0
mirror of https://github.com/deepspeedai/DeepSpeed.git synced 2025-10-20 15:33:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

Enable quantizer op on ROCm (#4114)

Browse Source 
This PR is to enable quantizer extension on ROCm

cc: @jithunnair-amd

---------

Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
This commit is contained in:
Ramya Ramineni
2024-01-10 17:48:54 -06:00
committed by GitHub
parent 6112d08ca7
commit f2d4845f07
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
op_builder/quantizer.py
View File

@ -32,4 +32,7 @@ class QuantizerBuilder(CUDAOpBuilder):
return ['csrc/includes'] return ['csrc/includes']
def extra_ldflags(self): def extra_ldflags(self):
return ['-lcurand'] if not self.is_rocm_pytorch():
return ['-lcurand']
else:
return []