Why does the watermark_remove_module.pt model file lack the ema key? #9
Description
Thanks for your excellent work. I am training of a model using the watermark_remove_module.pt file. I encountered an issue where the training process fails because the model file does not contain the ema key. It just contains 'ema unet', 'unet'.
Why does the watermark_remove_module.pt model file lack the ema key? Is this expected behavior or am I missing something?
Should I be using a different model for my task? If so, Which t2v model should I choose?
If the ema key is essential, how can I obtain a model file that includes it?
Thanks.
[rank0]: Traceback (most recent call last):
[rank0]: File "/home/VideoBooth/train_stage1.py", line 678, in
[rank0]: main(args)
[rank0]: File "/home/VideoBooth/train_stage1.py", line 395, in main
[rank0]: state_dict = torch.load(args.pretrained_t2v_model, map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()))["ema"]
[rank0]: KeyError: 'ema'
E0701 16:31:35.387000 139983327924608 torch/distributed/elastic/multiprocessing/api.py:826] failed (exitcode: 1) local_rank: 0 (pid: 1043993) of binary: /root/anaconda3/envs/videobooth/bin/python