# Empty lists to store the layers and the weights
all_layers = []; conv_weights = []

# Iterate over the model's structure
# (First level nn.Sequential)
for _, layer in enumerate(list(model.children())[0]):
    if type(layer) == nn.Conv2d:
        all_layers.append(layer)
        conv_weights.append(layer.weight)
    elif type(layer) in [nn.ReLU, nn.MaxPool2d]:
        all_layers.append(layer)
        conv_weights.append("*")

# Print layers and dimensions of weights
for idx, layer in enumerate(all_layers):
    print(f"{idx+1}. Layer: {layer}")
    if type(layer) == nn.Conv2d:
        print(f"          weights: {conv_weights[idx].shape}")
    else:
        print(f"          weights: {conv_weights[idx]}")
    print()