code: @Override
      public void onNDArrayPartial(INDArray arr, long idx, int... dimensions) {
          INDArray get = this.arr.get();
          get.tensorAlongDimension((int) idx, dimensions).assign(arr);
  
      }
docs: Used for partial updates using tensor along
  dimension
  @param arr        the array to count as an update
  @param idx        the index for the tensor along dimension
  @param dimensions the dimensions to act on for the tensor along dimension