Fix rlhf on trl v0.11.0

[satyaog]

Tested on:

remote [stdout] =================                                                                                                                                    
remote [stdout] Benchmark results                                                                                                                                    
remote [stdout] =================                                                                                                                                    
remote [stdout]                                                                                                                                                      
remote [stdout] System                                                                                                                                               
remote [stdout] ------                                                                                                                                               
remote [stdout] cpu:      AMD EPYC 7543 32-Core Processor                                                                                                            
remote [stdout] n_cpu:    64                                                                                                                                         
remote [stdout] product:  NVIDIA A100-SXM4-80GB                                                                                                                      
remote [stdout] n_gpu:    1                                                                                                                                          
remote [stdout] memory:   81920.0                                                                                                                                    
remote [stdout]                                                                                                                                                      
remote [stdout] Breakdown                                                                                                                                            
remote [stdout] ---------                                                                                                                                            
remote [stdout] bench       | fail |   n | ngpu |       perf |   sem% |   std% | peak_memory |      score | weight                                                   
remote [stdout] rlhf-single |    0 |   1 |    1 |    2589.65 |   0.5% |   4.2% |       12455 |    2589.65 |   1.00                                                   
remote [stdout]                                                                                                                                                      
remote [stdout] Scores                                                                                                                                               
remote [stdout] ------                                                                                                                                               
remote [stdout] Failure rate:       0.00% (PASS)                                                                                                                     
remote [stdout] Score:            2589.65                                                                                                                            

=================
Benchmark results
=================

System
------
cpu:      AMD EPYC 7543 32-Core Processor
n_cpu:    64
product:  NVIDIA A100-SXM4-80GB
n_gpu:    4
memory:   81920.0

Breakdown
---------
bench     | fail |   n | ngpu |       perf |   sem% |   std% | peak_memory |      score | weight
rlhf-gpus |    0 |   1 |    4 |    7179.24 |   0.4% |   2.9% |       21351 |    7179.24 |   1.00

Scores
------
Failure rate:       0.00% (PASS)
Score:            7179.24