TypeError: gather_nd() got an unexpected keyword argument 'batch_dims'

[nvnvashisth]

Any idea about this below error. Try to execute this.

 Instructions for updating:
 Use keras.layers.dense instead.
 Traceback (most recent call last):
   File "language_generation.py", line 686, in <module>
     main()
   File "language_generation.py", line 591, in main
     predictions, features = prediction_graph()
   File "language_generation.py", line 518, in prediction_graph_no_memory
     inp, inp_mask, seg_id, perm_mask, prev_tokens, prev_conf)
   File "language_generation.py", line 504, in body
     sampled_tokens, confidences = sample_token(logits)
   File "language_generation.py", line 276, in sample_token
     confidences = tf.gather_nd(params=probs, batch_dims=0, indices=samples)
 TypeError: gather_nd() got an unexpected keyword argument 'batch_dims'

I tried removing batch_dims argument and then it fails again after prompt. Any idea about this?

----PROMPT----

WARNING:tensorflow:From /Users/nitin/opt/anaconda3/envs/xlnet/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py:429: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.
Instructions for updating:
tf.py_func is deprecated in TF V2. Instead, use
    tf.py_function, which takes a python function which manipulates tf eager
    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to
    an ndarray (just call tensor.numpy()) but having access to eager tensors
    means `tf.py_function`s can use accelerators such as GPUs as well as
    being differentiable using a gradient tape.
    
  0%|                                                                                                  | 0/1 [00:00<?, ?it/s]
2019-12-23 16:59:50.810097: W tensorflow/core/framework/op_kernel.cc:1401] OP_REQUIRES failed at gather_nd_op.cc:50 : Invalid argument: indices[0] = [4712] does not index into param shape [1,32000]

Looking forward hearing from you.

[Omnis23]

Hi! I had a similar problem when running XLNet initially, but instead of using tf.gather_nd I used tf.gather, and in place of batch_dims=1 I used axis=1 and it is running great now. Hopefully this can apply to your situation as well!