Removed call to function_lib's register method

PiperOrigin-RevId: 481941823
google-deepmind · Oct 18, 2022 · 6b90fdb · 6b90fdb
1 parent 7298417
commit 6b90fdb
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/sonnet/src/recurrent.py b/sonnet/src/recurrent.py
@@ -1035,7 +1035,9 @@ def wrapper(*args, **kwargs):
               "api_implements": unique_api_name,
               "api_preferred_device": device
           })
-      function_lib.register(functions[device], *args, **kwargs)
+      concrete_func = functions[device].get_concrete_function(*args, **kwargs)
+      concrete_func.add_to_graph()
+      concrete_func.add_gradient_functions_to_graph()
     return functions[default](*args, **kwargs)
 
   return wrapper