diff --git a/entry.py b/entry.py
index c0726986..61f8eda3 100644
--- a/entry.py
+++ b/entry.py
@@ -118,7 +118,7 @@ torch.cuda.empty_cache()
 torch.cuda.ipc_collect()
 
 shape = (1, 4, 128, 128)
-randn = torch.randn(shape).cuda()
+randn = torch.randn(shape).to(torch.float16).cuda()
 
 
 def denoiser(input, sigma, c):
diff --git a/sgm/modules/diffusionmodules/openaimodel.py b/sgm/modules/diffusionmodules/openaimodel.py
index e19b83f9..ce16fd21 100644
--- a/sgm/modules/diffusionmodules/openaimodel.py
+++ b/sgm/modules/diffusionmodules/openaimodel.py
@@ -980,7 +980,7 @@ class UNetModel(nn.Module):
             self.num_classes is not None
         ), "must specify y if and only if the model is class-conditional"
         hs = []
-        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x)
         emb = self.time_embed(t_emb)
 
         if self.num_classes is not None: