char-t:

batch_hr = batch["hr"].to(opt.local_rank)                              
batch_lr = batch["lr"].to(opt.local_rank)                    
                                                                                                                                                 
g_hr = generator(batch_lr)                            
                                                                                                                                                                                                                               
d_real = discriminator(batch_hr)                                                          
d_fake = discriminator(g_hr.detach()) # detached from generator graph                
                                                                                                                                                       
loss_D = D(d_real, d_fake)                                    
                                                                                                                                                         
optimizer_D.zero_grad(set_to_none=True)                                                
loss_D.backward(retain_graph=True)                                                     
optimizer_D.step()                                                                     
                                                                                                                                                                                                                                      
loss_G1 = G1(batch_hr, g_hr)                                           
                                                                                                                                                  
d_real = discrim(batch_hr).detach() # detached from discrim graph                   
d_fake = discrim(g_hr)                                                         
                                                                                                                   
loss_G2 = G2(d_real, d_fake)                               
                                                                                                                                                       
loss_G = loss_G1 + loss_G2                                                                                                                           
                                                                     
optimizer_G.zero_grad(set_to_none=True)                                                    
loss_G.backward()                                                                          
optimizer_G.step()      



I think I’ve actually figured this out in the end! My focus on reducing the number of discriminator calls turned out to be a bit misplaced as it was definitely not the slow part of the training iteration. I’ve now got it setup where the discriminator trains and updates before creating new outputs which the generator uses for it’s loss function - basically splitting the code into clear train D, train G blocks as I probably should have done from the start