From 6c20e9d8981c145af7396bfcd38f24d4651e95f2 Mon Sep 17 00:00:00 2001
From: Avnish Narayan <38871737+avnishn@users.noreply.github.com>
Date: Wed, 16 Mar 2022 01:15:59 -0700
Subject: [PATCH] [RLlib] Change the slateq regression learning test with GPU
 to use torch only (#23168)

---
 release/rllib_tests/learning_tests/hard_learning_tests.yaml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/release/rllib_tests/learning_tests/hard_learning_tests.yaml b/release/rllib_tests/learning_tests/hard_learning_tests.yaml
index 6a64e2543..fac6d307d 100644
--- a/release/rllib_tests/learning_tests/hard_learning_tests.yaml
+++ b/release/rllib_tests/learning_tests/hard_learning_tests.yaml
@@ -468,6 +468,11 @@ slateq-interest-evolution-recsim-env:
         env_config:
             # Env class specified above takes one `config` arg in its c'tor:
             config:
+                # GPU support currently does not work for tensorflow. For the learning
+                # test for now, we'll use pytorch, and investigate why tensorflow is not working.
+                # One indication of the issue is that the loss inputs are not flattened from the
+                # original dictionary spaces.
+                framework: torch
                 # Each step, sample `num_candidates` documents using the env-internal
                 # document sampler model (a logic that creates n documents to select
                 # the slate from).