From 2ab8e6ff7e8565e83c85d85311e8d9d6a82d9b5f Mon Sep 17 00:00:00 2001 From: Pavel Sountsov Date: Wed, 1 Mar 2023 16:06:01 -0800 Subject: [PATCH] Use tfp.experimental.auto_composite_tensor rather than tfp.experimental.register_composite. PiperOrigin-RevId: 513368020 --- ibc/agents/ibc_policy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ibc/agents/ibc_policy.py b/ibc/agents/ibc_policy.py index 4371a23..72e98f4 100644 --- a/ibc/agents/ibc_policy.py +++ b/ibc/agents/ibc_policy.py @@ -65,7 +65,7 @@ def _create_variables(specs, training, step_type, network_state): return _create_variables -@tfp.experimental.register_composite +@tfp.experimental.auto_composite_tensor class MappedCategorical(tfp.distributions.Categorical): """Categorical distribution that maps classes to specific values."""