Type error in neural network with LSTM and a simple custom layer

[dkgaraujo]

Very interesting library, but it failed with my model. See below also the model's architecture. Note there is one custom layer, but it is only a layer that rescales its inputs before applying a softmax (as a temperature parameter would).

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[69], line 1
----> 1 img = visualtorch.lenet_view(model_arch, input_shape=input_shape)
      3 plt.axis("off")
      4 plt.tight_layout()

File ~/Coding/envs/.venv_mlx/lib/python3.12/site-packages/visualtorch/lenet_style.py:100, in lenet_view(model, input_shape, to_file, min_z, min_xy, max_xy, scale_z, scale_xy, type_ignore, index_ignore, color_map, one_dim_orientation, background_fill, padding, spacing, draw_funnel, shade_step, font, font_color, opacity, max_channels, offset_z)
     97 if color_map is not None:
     98     _color_map = defaultdict(dict, color_map)
--> 100 dummy_input = torch.rand(*input_shape)
    102 # Get the list of layers
    104 layers: OrderedDict[str, Any] = OrderedDict()

TypeError: rand() received an invalid combination of arguments - got (tuple, tuple), but expected one of:
 * (tuple of ints size, *, torch.Generator generator, tuple of names names, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (tuple of ints size, *, torch.Generator generator, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (tuple of ints size, *, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (tuple of ints size, *, tuple of names names, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)