generator.load_state_dict(checkpoint[ generator ]) KeyError: generator

It seems that the key "generator" is not present in the checkpoint dictionary. Make sure that you have saved the generator with the key "generator" while creating the checkpoint.

For example, if you are using PyTorch, you can save the generator like this:

torch.save({
    'generator': generator.state_dict(),
    'optimizer': optimizer.state_dict()
}, checkpoint_path)

Then, while loading the checkpoint, make sure to provide the same key to load the generator:

checkpoint = torch.load(checkpoint_path)
generator.load_state_dict(checkpoint['generator'])
optimizer.load_state_dict(checkpoint['optimizer'])

If the issue persists, try printing the keys in the checkpoint dictionary to see the available keys:

print(checkpoint.keys())

This will give you an idea if the key "generator" is present in the dictionary or not.