From 4c6351ef8ff3f4045a9a24bee6a994667d89c69c Mon Sep 17 00:00:00 2001
From: Acme Contributor <noreply@google.com>
Date: Wed, 3 Jan 2024 08:37:54 -0800
Subject: [PATCH] Added pytype None checks to inference_server.py

PiperOrigin-RevId: 595410959
Change-Id: Idf1c7d46fb3dfaccbc0e78c51cd91267f340a1e9
---
 acme/jax/inference_server.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/acme/jax/inference_server.py b/acme/jax/inference_server.py
index 9e4f27a684..e87071e6f3 100644
--- 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.