ZeroDivsion error?

Hi developer,

Thanks for sharing this work. I am trying to reproduce the MeshGraphNet developed by Deepminds in 2021, by using the module.EncodeProcessDecode and module.InteractionNetwroks.

I am using demo.tf2 as prototype and replace self._core with module.InteractionNetwork. I am only using test.tfrecord in original cylinder_flow case for easy testing the model on colab. (There is no global features, and num_nodes/num_edges equal to num_vertices/num_edges in that moment of mesh-space). I define each individual timestep as a ```GraphTuple```. Everything works fine until I started the actual training loop, then I got this error message:

StagingError: in user code:

File "", line 19, in update_step *
pred = model(inputs_tr, stats_list[0], stats_list[1], stats_list[2], stats_list[3], params.message_passing_steps)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "", line 98, in __call__ *
latent = self._encoder(inputs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "", line 31, in __call__ *
return self._network(inputs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/graph_nets/_base.py", line 78, in __call__ *
return self._build(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/graph_nets/modules.py", line 368, in _build *
globals=self._global_model(graph.globals))
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/graph_nets/_base.py", line 78, in __call__ *
return self._build(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/graph_nets/_base.py", line 112, in _build *
return self._model(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/sequential.py", line 68, in __call__ *
outputs = mod(outputs, *args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/nets/mlp.py", line 97, in __call__ *
inputs = layer(inputs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/linear.py", line 83, in __call__ *
self._initialize(inputs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/once.py", line 89, in wrapper *
_check_no_output(wrapped(*args, **kwargs))
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/utils.py", line 85, in _decorate_unbound_method *
return decorator_fn(bound_method, self, args, kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/base.py", line 262, in wrap_with_name_scope *
return method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/sonnet/src/linear.py", line 71, in _initialize *
stddev = 1 / math.sqrt(self.input_size)

ZeroDivisionError: float division by zero

I am not sure where I should look at to solve this. Any ideas?

Best,
ZY