From 4e990b4327d1234f0049d2eef323f4f5e7f97a2e Mon Sep 17 00:00:00 2001 From: Sergey Date: Wed, 29 Mar 2023 13:20:18 +0300 Subject: [PATCH] Update diffusion_decoder.py Chanhged import autocast from torch --- tortoise/models/diffusion_decoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tortoise/models/diffusion_decoder.py b/tortoise/models/diffusion_decoder.py index f67d21a..9d66e87 100644 --- a/tortoise/models/diffusion_decoder.py +++ b/tortoise/models/diffusion_decoder.py @@ -5,7 +5,7 @@ from abc import abstractmethod import torch import torch.nn as nn import torch.nn.functional as F -from torch import autocast +from torch.cuda.amp import autocast from tortoise.models.arch_util import normalization, AttentionBlock