update sampling code

ReadyToDance · Apr 5, 2023 · 72cd021 · 72cd021
1 parent 346ba99
commit 72cd021
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/lvdm/models/ddpm3d.py b/lvdm/models/ddpm3d.py
@@ -11,11 +11,10 @@
 
 import torch
 import torch.nn as nn
-from torch.optim.lr_scheduler import LambdaLR
-from torchvision.utils import make_grid
 import pytorch_lightning as pl
-from pytorch_lightning.utilities.distributed import rank_zero_only
-
+from torchvision.utils import make_grid
+from torch.optim.lr_scheduler import LambdaLR
+from pytorch_lightning.utilities import rank_zero_only
 from lvdm.models.modules.distributions import normal_kl, DiagonalGaussianDistribution
 from lvdm.models.modules.util import make_beta_schedule, extract_into_tensor, noise_like
 from lvdm.models.modules.lora import inject_trainable_lora