Added pytype None checks to inference_server.py

PiperOrigin-RevId: 595410959 Change-Id: Idf1c7d46fb3dfaccbc0e78c51cd91267f340a1e9
google-deepmind · Jan 3, 2024 · 4c6351e · 4c6351e
1 parent 1177501
commit 4c6351e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/acme/jax/inference_server.py b/acme/jax/inference_server.py
@@ -98,6 +98,9 @@ def _dereference_params(self, arg):
           key=self._keys,
           update_period=self._config.update_period)
 
+    if self._variable_client is None:
+      raise ValueError('_variable_client not set')
+
     params = self._variable_client.params
     device_idx = self._call_cnt % len(self._devices)
     # Select device via round robin, and update its params if they changed.