@ Override public void receiveUpdate ( INDArray array ) { extCounter . getAndIncrement ( ) ; updatesLock . writeLock ( ) . lock ( ) ; if ( updates == null ) { try ( MemoryWorkspace workspace = Nd4j . getMemoryManager ( ) . scopeOutOfWorkspaces ( ) ) { // TODO: this one has to be HOST-only if P2P is NOT supported updates = Nd4j . create ( array . shape ( ) , array . ordering ( ) ) ; } } hasSomething . compareAndSet ( false , true ) ; // if P2P is NOT supported - this call should be handled with cpu updates . addi ( array ) ; // we have to ensure, all operations were finished here Nd4j . getExecutioner ( ) . commit ( ) ; updatesLock . writeLock ( ) . unlock ( ) ; }
This method accepts updates suitable for StepFunction and puts them to the queue which is used in backpropagation loop