# choose layer to plot
target_layer = 0 #0 is input, 1 is 1st hidden etc