fix clip skip
This commit is contained in:
parent
98f9a7302d
commit
c816e25773
2
train.py
2
train.py
|
@ -636,7 +636,7 @@ def main(args):
|
||||||
encoder_hidden_states = text_encoder(cuda_caption, output_hidden_states=True)
|
encoder_hidden_states = text_encoder(cuda_caption, output_hidden_states=True)
|
||||||
|
|
||||||
if args.clip_skip > 0:
|
if args.clip_skip > 0:
|
||||||
encoder_hidden_states = encoder_hidden_states.hidden_states[-args.clip_skip]
|
encoder_hidden_states = text_encoder.text_model.final_layer_norm(encoder_hidden_states.hidden_states[-args.clip_skip])
|
||||||
else:
|
else:
|
||||||
encoder_hidden_states = encoder_hidden_states.last_hidden_state
|
encoder_hidden_states = encoder_hidden_states.last_hidden_state
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue