diff --git a/README.md b/README.md
index 2ae0c83..0b973c1 100644
--- a/README.md
+++ b/README.md
@@ -11,7 +11,7 @@ This is a fast, minimal port of Boris Dayma's [DALL·E Mega](https://github.com/
 To generate a 4x4 grid of DALL·E Mega images it takes:
 - 89 sec with a T4 in Colab
 - 48 sec with a P100 in Colab
-- 14 sec with an A100 on Replicate
+- 13 sec with an A100 on Replicate
 
 The flax model and code for converting it to torch can be found [here](https://github.com/kuprel/min-dalle-flax).
 
diff --git a/cog.yaml b/cog.yaml
index 9d2e90b..bed81dc 100644
--- a/cog.yaml
+++ b/cog.yaml
@@ -1,13 +1,13 @@
 build:
-  cuda: "11.4"
+  cuda: "11.5.1"
   gpu: true
   python_version: "3.10"
   system_packages:
     - "libgl1-mesa-glx"
     - "libglib2.0-0"
   python_packages:
-    - "min-dalle==0.3.5"
+    - "min-dalle==0.3.7"
   run:
-    - pip install torch==1.11.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html
+    - pip install torch==1.12.0+cu116 -f https://download.pytorch.org/whl/torch_stable.html
 
-predict: "replicate_predictor.py:ReplicatePredictor"
+predict: "replicate_predictor.py:ReplicatePredictor"
\ No newline at end of file
diff --git a/min_dalle/min_dalle.py b/min_dalle/min_dalle.py
index fedbf3a..9b02bba 100644
--- a/min_dalle/min_dalle.py
+++ b/min_dalle/min_dalle.py
@@ -4,15 +4,21 @@ import numpy
 from torch import LongTensor, FloatTensor
 from math import sqrt
 import torch
+import torch.backends.cudnn, torch.backends.cuda
 import json
 import requests
 from typing import Iterator
-torch.set_grad_enabled(False)
-torch.set_num_threads(os.cpu_count())
-
 from .text_tokenizer import TextTokenizer
 from .models import DalleBartEncoder, DalleBartDecoder, VQGanDetokenizer
 
+torch.set_grad_enabled(False)
+torch.set_num_threads(os.cpu_count())
+torch.backends.cudnn.enabled = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True
+
 MIN_DALLE_REPO = 'https://huggingface.co/kuprel/min-dalle/resolve/main/'
 
 
diff --git a/min_dalle/models/dalle_bart_decoder.py b/min_dalle/models/dalle_bart_decoder.py
index 917f354..9e75e4f 100644
--- a/min_dalle/models/dalle_bart_decoder.py
+++ b/min_dalle/models/dalle_bart_decoder.py
@@ -1,8 +1,6 @@
 from typing import Tuple, List
 import torch
 from torch import nn, LongTensor, FloatTensor, BoolTensor
-torch.set_grad_enabled(False)
-
 from .dalle_bart_encoder import GLU, AttentionBase
 
 IMAGE_TOKEN_COUNT = 256
@@ -180,6 +178,7 @@ class DalleBartDecoder(nn.Module):
             self.zero_prob,
             torch.exp(logits - top_logits[:, [0]])
         )
+        probs[:, 2 ** 14:] = 0              # vqgan vocab_count is only 2 ** 14
         return probs, attention_state
 
 
diff --git a/min_dalle/models/dalle_bart_encoder.py b/min_dalle/models/dalle_bart_encoder.py
index 7081fa3..cb93389 100644
--- a/min_dalle/models/dalle_bart_encoder.py
+++ b/min_dalle/models/dalle_bart_encoder.py
@@ -1,7 +1,6 @@
 from typing import List
 import torch
 from torch import nn, BoolTensor, FloatTensor, LongTensor
-torch.set_grad_enabled(False)
 
 
 class GLU(nn.Module):
diff --git a/min_dalle/models/vqgan_detokenizer.py b/min_dalle/models/vqgan_detokenizer.py
index c3ffebb..9c9ee48 100644
--- a/min_dalle/models/vqgan_detokenizer.py
+++ b/min_dalle/models/vqgan_detokenizer.py
@@ -1,7 +1,6 @@
 import torch
 from torch import FloatTensor, LongTensor
 from torch.nn import Module, ModuleList, GroupNorm, Conv2d, Embedding
-torch.set_grad_enabled(False)
 
 
 class ResnetBlock(Module):
diff --git a/replicate_predictor.py b/replicate_predictor.py
index 5177c5d..39018ec 100644
--- a/replicate_predictor.py
+++ b/replicate_predictor.py
@@ -1,9 +1,11 @@
 from min_dalle import MinDalle
 import tempfile
-import torch
+import torch, torch.backends.cudnn
 from typing import Iterator
 from cog import BasePredictor, Path, Input
 
+torch.backends.cudnn.deterministic = False
+
 
 class ReplicatePredictor(BasePredictor):
     def setup(self):
diff --git a/setup.py b/setup.py
index 0c8b658..b408af3 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ setuptools.setup(
     name='min-dalle',
     description = 'min(DALL·E)',
     # long_description=(Path(__file__).parent / "README.rst").read_text(),
-    version='0.3.5',
+    version='0.3.7',
     author='Brett Kuprel',
     author_email='brkuprel@gmail.com',
     url='https://github.com/kuprel/min-dalle',