diff --git a/stacks/pinot-noir/packages_jed.yaml b/stacks/pinot-noir/packages_jed.yaml index dcc04e3ecacc5be4b2e480c5b8056bd880dae8f8..c260f3e315b6c8ae89ed64c00531d8640f2eee5d 100644 --- a/stacks/pinot-noir/packages_jed.yaml +++ b/stacks/pinot-noir/packages_jed.yaml @@ -15,12 +15,12 @@ packages: # --------------------------------------------------------------------------- hypre: require: - - spec: '+cuda cuda_arch=90' + - spec: '+cuda cuda_arch=90 +unified-memory' when: '%gcc' kokkos: require: - - spec: '+cuda cuda_arch=90 +cuda_uvm' + - spec: '+cuda cuda_arch=90 +cuda_uvm +wrapper' when: '%gcc' openmpi: @@ -32,6 +32,10 @@ packages: - spec: '+cuda cuda_arch=90' when: '%gcc' + py-tensorflow: + require: + - spec: '+cuda cuda_arch=90 +nccl' + py-torch: require: - spec: '+cuda cuda_arch=90 +nccl +cudnn' @@ -46,7 +50,7 @@ packages: ucx: require: - - spec: '+xpmem +cma +rdmacm +rc +ud +dc +verbs +ib_hw_tm +vfs +gdrcopy +cuda cuda_arch=90 target=icelake ~rocm' + - spec: '+xpmem +cma +rdmacm +rc +ud +dc +verbs +ib_hw_tm +gdrcopy +cuda cuda_arch=90 target=icelake ~rocm' # --------------------------------------------------------------------------- # Externals