fix(modeling_ilql): single q-head indexing (#471)

* fix(modeling_ilql): single q-head indexing in `.generate`

* fix(modeling_ilql): same error for seq2seq

---------

Co-authored-by: reciprocated <[email protected]>

Author: maxreciprocate

trlx/models/modeling_ilql.py

@@ -289,7 +289,7 @@ def generate(
             if self.two_qs:
                 qs = torch.minimum(target_qs[0][:, -1, :], target_qs[1][:, -1, :])
             else:
-                qs = target_qs[:, -1, :]
+                qs = target_qs[0][:, -1, :]
 
             logits = logits[:, -1, :]
             vs = vs[:, -1, :]
@@ -469,7 +469,7 @@ def generate(
             if self.two_qs:
                 qs = torch.minimum(target_qs[0][:, -1, :], target_qs[1][:, -1, :])
             else:
-                qs = target_qs[:, -1, :]
+                qs = target_qs[0][:, -1, :]
 
             logits = logits[:, -1, :]
             vs = vs[:, -1, :]