Chesscorner commited on
Commit
4777031
·
verified ·
1 Parent(s): 19b2caf

Update caption.py

Browse files
Files changed (1) hide show
  1. caption.py +2 -2
caption.py CHANGED
@@ -9,8 +9,8 @@ processor = AutoProcessor.from_pretrained("Chesscorner/git-chess-v3")
9
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10
  model.to(device)
11
 
12
- max_length = 50
13
- num_beams = 4
14
  gen_kwargs = {'max_length': max_length, 'num_beams': num_beams}
15
 
16
  def predict_step(image):
 
9
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10
  model.to(device)
11
 
12
+ max_length = 60
13
+ num_beams = 2
14
  gen_kwargs = {'max_length': max_length, 'num_beams': num_beams}
15
 
16
  def predict_step(image):