Better behavior if something goes wrong

torchexplorer/hook/hook.py

@@ -191,6 +191,7 @@ def record_sizes(tensors: tuple[OTensor, ...], tensor_trackers: list[SizeTracker
  assert stored_shape is not None
  if len(shape) != len(stored_shape):
  tensor_trackers[i].size = None
+ continue
 
  # TODO: with graphviz caching, this actually doesn't need to run
  # since only the first pass sizes are stored.
@@ -336,6 +337,8 @@ def process_tensor(tensor: Tensor) -> Optional[GradFn]:
  return tuple(process_tensor(tensor) for tensor in utils.iter_not_none(tensors))
 
 def _get_next_gradfns(tensors: tuple[OTensor, ...]) -> tuple[Optional[GradFn], ...]:
+ if len(tensors) == 0:
+ return tuple()
  # Hacky workaround, couldn't figure out import
  # Check if all gradfns are the same 'BackwardHookFunctionBackward'
  backhook_class = 'BackwardHookFunctionBackward'

torchexplorer/render/layout.py

@@ -206,8 +206,11 @@ def _translate_inner_layouts(layout: NodeLayout) -> None:
  if is_input_node(l.display_name):
  input_centers.append(_center(l))
 
- center = np.mean(np.array(input_centers), axis=0)
- trans = utils.list_add(target_input_pos, [-center[0], -center[1]])
+ if len(input_centers) > 0:
+ center = np.mean(np.array(input_centers), axis=0)
+ trans = utils.list_add(target_input_pos, [-center[0], -center[1]])
+ else:
+ trans = target_input_pos
 
  def apply_translation(l: Union[NodeLayout, TooltipLayout]):
  l.bottom_left_corner = utils.list_add(l.bottom_left_corner, trans)