lihaoxin2020 commited on
Commit
61d0002
1 Parent(s): 364e345

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -28,6 +28,9 @@ YT_LENGTH_LIMIT_S = 3600 # limit to 1 hour YouTube files
28
  # device = "cpu"
29
  # model_kwargs = {}
30
  device = "cuda"
 
 
 
31
  # define the pipeline
32
  pipe = pipeline(
33
  model=MODEL_NAME,
 
28
  # device = "cpu"
29
  # model_kwargs = {}
30
  device = "cuda"
31
+ torch_dtype = torch.bfloat16
32
+ model_kwargs = {'attn_implementation': 'sdpa'}
33
+
34
  # define the pipeline
35
  pipe = pipeline(
36
  model=MODEL_NAME,