diff --git a/scripts/tilevae.py b/scripts/tilevae.py
index 199eb7c..5f4ed4a 100644
--- a/scripts/tilevae.py
+++ b/scripts/tilevae.py
@@ -234,7 +234,7 @@ def custom_group_norm(input, num_groups, mean, var, weight=None, bias=None, eps=
     input_reshaped = input.contiguous().view(
         1, int(b * num_groups), channel_in_group, *input.size()[2:])
 
-    out = F.batch_norm(input_reshaped, mean, var, weight=None, bias=None, training=False, momentum=0, eps=eps)
+    out = F.batch_norm(input_reshaped, mean.to(input), var.to(input), weight=None, bias=None, training=False, momentum=0, eps=eps)
     out = out.view(b, c, *input.size()[2:])
 
     # post affine transform