JMalott commited on
Commit
aa8e101
1 Parent(s): e38321c

Update min_dalle/min_dalle.py

Browse files
Files changed (1) hide show
  1. min_dalle/min_dalle.py +2 -2
min_dalle/min_dalle.py CHANGED
@@ -12,8 +12,8 @@ from .models import DalleBartEncoder, DalleBartDecoder, VQGanDetokenizer
12
  import streamlit as st
13
  import time
14
 
15
- torch.set_grad_enabled(True)
16
- torch.set_num_threads(1)
17
  torch.backends.cudnn.enabled = True
18
  torch.backends.cudnn.allow_tf32 = True
19
 
 
12
  import streamlit as st
13
  import time
14
 
15
+ torch.set_grad_enabled(False)
16
+ torch.set_num_threads(torch.set_num_threads(os.cpu_count()))
17
  torch.backends.cudnn.enabled = True
18
  torch.backends.cudnn.allow_tf32 = True
19