Torch.utils.tensorboard fails with multiple inputs model

Hi all! I have a matching network model which receives 4 inputs, while it all works correctly, once I add in the tensorboard add_graph functioni n the mix it fails miserably.

This is how the model is called:
acc, c_loss, preds = model(x_support_set, y_support_set_one_hot, x_target, y_target)

And this is the add_graph call:
writer.add_graph(model, (x_support_set, y_support_set_one_hot, x_target, y_target), True)

The error is from an internal method in the model where I encode the inputs.

Full trace here:

/home/davide/Desktop/coefficient/native_oneshot/native_oneshot/matching_network/model/matching.py:90: TracerWarning: Converting a tensor to a Python index might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  batch_size = len(target_images)
/home/davide/Desktop/coefficient/native_oneshot/native_oneshot/matching_network/model/matching.py:55: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  print(support_set_images.size(), target_images.size())
torch.Size([1, 25, 3, 100, 100]) torch.Size([10, 3, 100, 100])
/home/davide/Desktop/coefficient/native_oneshot/native_oneshot/matching_network/model/matching.py:56: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  for i in np.arange(support_set_images.size(1)):
Integer division of tensors using div or / is no longer supported, and in a future release div will perform true division as in Python 3. Use true_divide or floor_divide (// in Python) instead.
Error occurs, No graph saved
 96%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍         | 24/25 [00:06<00:00,  3.85it/s]
Traceback (most recent call last):
  File "/home/davide/.pyenv/versions/3.7.8/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/home/davide/.pyenv/versions/3.7.8/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/davide/Desktop/coefficient/native_oneshot/native_oneshot/scripts/matchnet_train.py", line 298, in <module>
    log_dir=model_folder_path,
  File "/home/davide/Desktop/coefficient/native_oneshot/native_oneshot/scripts/matchnet_train.py", line 203, in run_epoch
    batch_number=i,
  File "/home/davide/Desktop/coefficient/native_oneshot/native_oneshot/scripts/matchnet_train.py", line 174, in run_batch
    True,
  File "/home/davide/.virtualenvs/native/lib/python3.7/site-packages/torch/utils/tensorboard/writer.py", line 714, in add_graph
    self._get_file_writer().add_graph(graph(model, input_to_model, verbose))
  File "/home/davide/.virtualenvs/native/lib/python3.7/site-packages/torch/utils/tensorboard/_pytorch_graph.py", line 291, in graph
    raise e
  File "/home/davide/.virtualenvs/native/lib/python3.7/site-packages/torch/utils/tensorboard/_pytorch_graph.py", line 285, in graph
    trace = torch.jit.trace(model, args)
  File "/home/davide/.virtualenvs/native/lib/python3.7/site-packages/torch/jit/__init__.py", line 955, in trace
    check_tolerance, strict, _force_outplace, _module_class)
  File "/home/davide/.virtualenvs/native/lib/python3.7/site-packages/torch/jit/__init__.py", line 1109, in trace_module
    module._c._create_method_from_trace(method_name, func, example_inputs, var_lookup_fn, strict, _force_outplace)
  File "/home/davide/.virtualenvs/native/lib/python3.7/site-packages/torch/nn/modules/module.py", line 720, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File "/home/davide/.virtualenvs/native/lib/python3.7/site-packages/torch/nn/modules/module.py", line 704, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/davide/Desktop/coefficient/native_oneshot/native_oneshot/matching_network/model/matching.py", line 77, in forward
    similarities = self.predict_similarities(support_set_images, target_images)
  File "/home/davide/Desktop/coefficient/native_oneshot/native_oneshot/matching_network/model/matching.py", line 91, in predict_similarities
    encodings = self.encode_support_and_target(support_set_images, target_images)
  File "/home/davide/Desktop/coefficient/native_oneshot/native_oneshot/matching_network/model/matching.py", line 56, in encode_support_and_target
    for i in np.arange(support_set_images.size(1)):
RuntimeError: Integer division of tensors using div or / is no longer supported, and in a future release div will perform true division as in Python 3. Use true_divide or floor_divide (// in Python) instead.

specifically the exception is thrown on this line for i in np.arange(support_set_images.size(1)): where support_set_images is of size torch.Size([1, 25, 3, 100, 100])

I’m a bit lost as to why everything would train fine, and when adding the add_graph call would throw such an exception deeper in the code, while the correct data is being passed. Any help would be appreciated!

As a side note, I already verified that the tensors are both FloatTensors as per this thread RuntimeError: Integer division of tensors using div or / is no longer supported, and in a future release div will perform true division as in Python 3. Use true_divide or floor_divide (// in Python) instead