Anyway a good thing you can do is check for NaNs in the loss (it’s a cheap operation).
If there is a NaN you can follow back the error and revise the input or the layers and try to find the reasoin reproducing that

Can you tell me how I’d be able to do that? So far, I’ve only found suggestions on using something called autograd.set_detect_anomaly(True) which I used as follows:
with torch.set_grad_enabled(phase == 'train'):
    with torch.autograd.set_detect_anomaly(True):
        audio_output = model_a(audio_features) # prediction
                                
        #audio_output = torch.nan_to_num(model_a(audio_features), neginf=0.0, posinf=1.0) # make prediction

        if (isinstance(audio_output,GoogLeNetOutputs)): # extract tensor from GoogLeNetOutputs
            audio_output = audio_output.logits

            preds =  torch.argmax(audio_output, 1) # get prediction
            # print(f"preds: {preds}, audio_output: {audio_output}, labels: {labels}\n")
            a_loss = criterion(audio_output, labels) # calculate loss
            print(f"A Loss: {a_loss}")
            #a_loss = a1_loss.float() # convert from long to float

            if phase == 'train':
                print('\n--- IF ---\n')
                a_loss.backward()

                for name, param in model_a.named_parameters():
                    print(name, torch.isfinite(param.grad).all()) # print grads

                model_a_optimizer.step()

                print('\n--- STEP ---\n')
                for name, param in model_a.named_parameters():
                    print(name, torch.isfinite(param.grad).all()) # print grads
                            
                #statistics
                print('a: ', a_loss, a_loss.item(), audio_features.size(0))
                a_running_loss += a_loss.item() * audio_features.size(0)
                a_running_corrects += torch.sum(preds == labels.data)
                                
                audio_dataset_sizes[phase] = audio_dataset_sizes[phase]+1

Should I be using set_detect_anomaly or should I use something else?

There are very few resons why you would get a NaN.

I assume you mean get a NaN in the loss component, correct?