How to visualize the architecture of a model?

Hello I am very new to this field and
dealing with Pytorch is difficult for me.

How can I visualize the architecture of a model?

I have tried the following:

and it throws an error that I cannot solve.

The error suggests your model is in training mode but should be in evaluation mode. So, can you put model.eval() before the yhat = model(batch) line, and see if that solves the error?

batch = next(iter(training_data))

Thank you for your reply. After your addition I received another error message.

What does “expecting the last two dimensions of the Tensor to be H and W instead got {img.shape[-2:]}” mean?
And how can I fulfill the expectation?

Your variable img isn’t a Tensor it’s a list, so check what img exactly is. It could be a list of a single tensor, but you’ll need to check. You could try printing len(img) to see how many items are in the list

You could also try torchview. It accepts any type of input and output for your model.
Here is the link: