I suggest to profile your code with then mentioned resources so you can investigate the bottlenecks instead of guessing.
Profile help:

  

      github.com/tensorflow/ranking
  

  
    
  
	  
  

  
    
      Slow training due to tf.gather 1
    

    
      
        opened 
        
          
        
        Jan 24, 2020
      
      

        
          closed 
        
          
        
        Feb 6, 2020
      
        

      
        
          
          eladBenedict
        
      
    

    
    
  


  
    Hi,

We, at Taboola, are working on implementing a learning-to-rank model, wit…h your framework, for our large scale production environment.

For small batch sizes (e.g. 32) the number of steps/sec during training is reasonable. However, it seems to grow significantly as the batch size grows. For example, processing a single batch of size 256 takes a few seconds(!).

Running profiling, it seems that the issue is in [tf.gather] (https://github.com/tensorflow/ranking/blob/a6fcfa2ab5a17e3e112dceb18e628ab1593385ba/tensorflow_ranking/python/model.py#L364)

Attached are:
1. The profiling json (renamed from .json to .txt due to upload restrictions)
2. A screenshot of the presumably problematic op in the profiling session

The profiling screenshot:
![image (5)](https://user-images.githubusercontent.com/11731485/73085519-7e88b280-3ed7-11ea-9377-c7cb76b1a4cb.png)

The profiling timeline:
[timeline-560-128.txt](https://github.com/tensorflow/ranking/files/4109481/timeline-560-128.txt)

Just for some context, in this scenario list_size=4 and group_size=1.

Is this a known issue?
Please advise.

Cheers, and thanks for the great framework!
Elad