Adapt env common tests

alexhernandezgarcia · Apr 2, 2024 · 8a7f6ab · 8a7f6ab
1 parent 51939bb
commit 8a7f6ab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/gflownet/envs/common.py b/tests/gflownet/envs/common.py
@@ -467,7 +467,6 @@ def test__gflownet_minimal_runs(self, n_repeat=1):
                 float_precision=config.float_precision,
                 base=forward_policy,
             )
-            self.env.proxy = proxy  # Set proxy in env.
             config.env.buffer.train = None  # No buffers
             config.env.buffer.test = None
             config.env.buffer.replay_capacity = 0  # No replay buffer
@@ -478,6 +477,7 @@ def test__gflownet_minimal_runs(self, n_repeat=1):
                 device=config.device,
                 float_precision=config.float_precision,
                 env_maker=self.env.__class__,
+                proxy=proxy,
                 forward_policy=forward_policy,
                 backward_policy=backward_policy,
                 buffer=config.env.buffer,