sayakpaul HF staff commited on
Commit
9435d99
1 Parent(s): 3304f7d

fix: state dict call.

Browse files
Files changed (1) hide show
  1. convert.py +1 -1
convert.py CHANGED
@@ -71,7 +71,7 @@ def run_conversion(text_encoder_weights: str = None, unet_weights: str = None):
71
  text_encoder_state_dict_from_pt = pt_text_encoder.state_dict()
72
  run_assertion(text_encoder_state_dict_from_pt, text_encoder_state_dict_from_tf)
73
  if unet_weights is None:
74
- unet_state_dict_from_pt = pt_text_encoder.state_dict()
75
  run_assertion(unet_state_dict_from_pt, unet_state_dict_from_tf)
76
 
77
  print(
 
71
  text_encoder_state_dict_from_pt = pt_text_encoder.state_dict()
72
  run_assertion(text_encoder_state_dict_from_pt, text_encoder_state_dict_from_tf)
73
  if unet_weights is None:
74
+ unet_state_dict_from_pt = pt_unet.state_dict()
75
  run_assertion(unet_state_dict_from_pt, unet_state_dict_from_tf)
76
 
77
  print(