From b599adc938de5a9d728dd73dc311ab95e1f786e0 Mon Sep 17 00:00:00 2001
From: Isotr0py <2037008807@qq.com>
Date: Wed, 8 Feb 2023 19:34:03 +0800
Subject: [PATCH] fix Input type error when using DDP

---
 train_network.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/train_network.py b/train_network.py
index c2f9cbf6..fc387bc3 100644
--- a/train_network.py
+++ b/train_network.py
@@ -1,5 +1,6 @@
 from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION
 from torch.optim import Optimizer
+from torch.cuda.amp import autocast
 from typing import Optional, Union
 import importlib
 import argparse