jax_predict.py formating

Author: pbenner

equitrain/backends/jax_predict.py

@@ -148,9 +148,7 @@ def _chunk_iterator():
 
     def _graph_real_counts(graph):
         try:
-            pad_mask = np.asarray(
-                jraph.get_graph_padding_mask(graph), dtype=bool
-            )
+            pad_mask = np.asarray(jraph.get_graph_padding_mask(graph), dtype=bool)
             n_node = np.asarray(graph.n_node)
 
             # Loader may append zero-node graphs to satisfy multi-device padding.