diff --git a/nd4j/nd4j-backends/nd4j-backend-impls/nd4j-cuda/src/main/java/org/nd4j/jita/handler/impl/CudaZeroHandler.java b/nd4j/nd4j-backends/nd4j-backend-impls/nd4j-cuda/src/main/java/org/nd4j/jita/handler/impl/CudaZeroHandler.java index 23301f4be..106ac9c3a 100644 --- a/nd4j/nd4j-backends/nd4j-backend-impls/nd4j-cuda/src/main/java/org/nd4j/jita/handler/impl/CudaZeroHandler.java +++ b/nd4j/nd4j-backends/nd4j-backend-impls/nd4j-cuda/src/main/java/org/nd4j/jita/handler/impl/CudaZeroHandler.java @@ -1184,9 +1184,10 @@ public class CudaZeroHandler implements MemoryHandler { protected synchronized cublasHandle_t getCudaCublasHandle(OpaqueLaunchContext lc) { val deviceId = Nd4j.getAffinityManager().getDeviceForCurrentThread(); - if (cublasHandles.get(deviceId) == null) + if (cublasHandles.get(deviceId) == null) { cublasHandles.remove(deviceId); cublasHandles.add(deviceId, new cublasHandle_t(nativeOps.lcBlasHandle(lc))); + } return cublasHandles.get(deviceId); }