Yes, you could use forward hooks to store the activation from a specific layer (I assume you meant activation instead of weight) as described here and use the linked umap library to reduce the dimensionality. Once this is done, you could use matplotlib to create a scatter plot (take a look at the examples at the bottom of the docs).