nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-04-13 06:46:08 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

debug

Browse source
This commit is contained in:
Shen-Chenhui 2024-04-12 18:35:28 +08:00
parent c3763cb556
commit 8aad3d6de2
1 changed files with 11 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
opensora/models/vae/vae_3d_v2.py
View file

@ -915,9 +915,19 @@ class VEALoss(nn.Module):
kl_loss_weight = 0.000001,
vgg=None,
vgg_weights: VGG16_Weights = VGG16_Weights.DEFAULT,
device = "cpu",
dtype = "bf16"
):
super().__init__()
if type(dtype) == str:
if dtype == "bf16":
dtype = torch.bfloat16
elif dtype == "fp16":
dtype = torch.float16
else:
raise NotImplementedError(f'dtype: {dtype}')
# KL Loss
self.kl_loss_weight = kl_loss_weight
# Perceptual Loss
@ -929,7 +939,7 @@ class VEALoss(nn.Module):
weights = vgg_weights
)
vgg.classifier = Sequential(*vgg.classifier[:-2])
self.vgg = vgg
self.vgg = vgg.to(device, dtype)
def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer):
breakpoint() # TODO: scrutinize