Attention Maps

I'd like to visualise the PERCH 'bird-vocalization-classifier' model's attention over a spectrogram however I'm unable to calculate the gradients directly. 
I'm new to TF but I think I would need either the JAX weights or TF weights with gradients enabled, which I can't find anywhere. Is it possible to access this? 
Many Thanks!