train 2 unet error

hello I want to try to train two unet this is my train code

def train(img):
    global LOSS
    Loss=[0,0]
    for i in range(2):
        name = "image_checkpoint.pth"
        trainer.load(name)
        loss = trainer(img,unet_number =i+1)
        Loss[i] = Loss[i]+loss
        trainer.update(unet_number=i+1)
        trainer.save(name)
    wandb.log({"u1_loss":Loss[0],"u2_loss":Loss[1]},step=step)
    LOSS[0]=LOSS[0]+Loss[0]
    LOSS[1]=LOSS[1]+Loss[1]

but an error occurs

you cannot only train on one unet at a time. you will need to save the trainer into a checkpoint, and resume training on a new unet

Therefore I changed my code

def train(img):
    global LOSS
    Loss=[0,0]
    for i in range(2):
        name = "image_checkpoint"+str(i+1)+".pth"
        trainer.load(name)
        loss = trainer(img,unet_number =i+1)
        Loss[i] = Loss[i]+loss
        trainer.update(unet_number=i+1)
        trainer.save(name)
    wandb.log({"u1_loss":Loss[0],"u2_loss":Loss[1]},step=step)
    LOSS[0]=LOSS[0]+Loss[0]
    LOSS[1]=LOSS[1]+Loss[1]

But he still got the above error

How to train two unet ？ thank

[TheFusion21]

afaik you have to reinstantiate the trainer and/or imagen and then load the cp

[dinhdungz]

@TheFusion21 Thanks you, that error is gone