for batch_idx,img in enumerate(loaders['rgbimage_dataloader']):
        optimizer.zero_grad()
        # move to GPU
        if use_cuda:
            real = img[0].to(device)
            sketch_image = cycle_model(real)
        else:
            real = img[0].to(device)
            sketch_image = cycle_model(real)
        
        styleGan_output = model(real,sketch_image,Target_color.float(),noise)
       ####Pretrained model#################
        fakeEdge_image = cycle_model(styleGan_output)   
        
        Edge_loss = criterion(fakeEdge_image,sketch_image)   
        total_loss = Edge_loss
        total_loss.backward()
        optimizer.step()
        
       ####This is the way i am saving ########################
       torch.save(model.state_dict(),f"model_history/mu_stylegan{epoch}.pth")

In this way is the the model saving state of both the pretrained model as well ?
or How could we do that ?