docs: @Override
      public INDArray getWordVectorMatrix(String word) {
          // TODO: add variable UNK here
          int idx = 0;
          if (hasWord(word))
              idx = vocabCache.indexOf(word);
          else if (getUNK() != null)
              idx = vocabCache.indexOf(getUNK());
          else
              return null;
  
          int deviceId = Nd4j.getAffinityManager().getDeviceForCurrentThread();
          INDArray array = null;
  
          if (cachePerDevice > 0 && cacheWrtDevice.get(deviceId).containsKey(idx))
              return cacheWrtDevice.get(Nd4j.getAffinityManager().getDeviceForCurrentThread()).get(idx);
  
          array = storage.get(idx);
  
          if (cachePerDevice > 0) {
              // TODO: add cache here
              long arrayBytes = array.length() * array.data().getElementSize();
              if ((arrayBytes * cacheWrtDevice.get(deviceId).size()) + arrayBytes < cachePerDevice)
                  cacheWrtDevice.get(deviceId).put(idx, array);
          }
  
          return array;
      }