Fixed disabling of gradients in the torch code

2 years ago · a3a247e6ec
parent 9d6b6dcc92
commit a3a247e6ec
1 changed files with 1 additions and 1 deletions
--- a/min_dalle/min_dalle_torch.py
+++ b/min_dalle/min_dalle_torch.py
@ -2,7 +2,7 @@ import numpy
 from typing import Dict
 from torch import LongTensor, FloatTensor
 import torch
-torch.no_grad()
+torch.set_grad_enabled(False)

 from .models.vqgan_detokenizer import VQGanDetokenizer
 from .models.dalle_bart_encoder_torch import DalleBartEncoderTorch