[RLlib] Fix tfp import (ray-project#34756)

Signed-off-by: Rohan Potdar <[email protected]>
pcmoritz · May 24, 2023 · 09e07fd · 09e07fd
1 parent 806b633
commit 09e07fd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/rllib/algorithms/bandit/bandit_tf_model.py b/rllib/algorithms/bandit/bandit_tf_model.py
@@ -1,13 +1,13 @@
 import gymnasium as gym
-import tensorflow_probability as tfp
 
 from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
 from ray.rllib.utils.annotations import override
-from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.framework import try_import_tf, try_import_tfp
 from ray.rllib.utils.typing import TensorType
 
 tf1, tf, tfv = try_import_tf()
+tfp = try_import_tfp()
 
 
 class OnlineLinearRegression(tf.Module if tf else object):