Merge pull request #24 from TheFutureGadgetsLab/main

Fixed disabling of gradients in the torch code
2 years ago · b8c4173181
parent 9d6b6dcc92 a3a247e6ec
commit b8c4173181
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