Tonic commited on
Commit
b3758b8
1 Parent(s): 163f1eb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -3
app.py CHANGED
@@ -1,11 +1,10 @@
1
  import gradio as gr
2
  from transformers import pipeline
 
3
  import torch
4
 
5
  # Set max_split_size_mb
6
- torch.backends.cuda.matmul.allow_tf32 = True
7
- torch.backends.cudnn.allow_tf32 = True
8
- torch.backends.cuda.matmul.max_split_size_mb = 50 # You can adjust this value
9
 
10
  raven_pipeline = pipeline(
11
  "text-generation",
 
1
  import gradio as gr
2
  from transformers import pipeline
3
+ import os
4
  import torch
5
 
6
  # Set max_split_size_mb
7
+ os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:50'
 
 
8
 
9
  raven_pipeline = pipeline(
10
  "text-generation",