Chesscorner commited on
Commit
f63e78d
·
verified ·
1 Parent(s): ecab7e2

Update caption.py

Browse files
Files changed (1) hide show
  1. caption.py +3 -2
caption.py CHANGED
@@ -15,8 +15,9 @@ use_fp16 = device.type == "cuda"
15
  if use_fp16:
16
  model.half()
17
 
18
- # Set generation parameters
19
- gen_kwargs = {'max_length': 50, 'num_beams': 2} # Adjust num_beams if needed
 
20
 
21
 
22
  # Prediction function
 
15
  if use_fp16:
16
  model.half()
17
 
18
+ max_length = 16
19
+ num_beams = 4
20
+ gen_kwargs = {'max_length': max_length, 'num_beams': num_beams}
21
 
22
 
23
  # Prediction function