diff --git a/__init__.py b/__init__.py
index 93a3aba..f5eb101 100644
--- a/__init__.py
+++ b/__init__.py
@@ -40,14 +40,12 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
     else:
         if "noise_mask" in latent:
             noise_mask = latent['noise_mask']
-            noise_mask = torch.nn.functional.interpolate(noise_mask[None, None,], size=(noise.shape[2], noise.shape[3]),
-                                                         mode="bilinear")
+            noise_mask = torch.nn.functional.interpolate(noise_mask[None, None,], size=(noise.shape[2], noise.shape[3]), mode="bilinear")
             noise_mask = noise_mask.round()
             noise_mask = torch.cat([noise_mask] * noise.shape[1], dim=1)
             noise_mask = torch.cat([noise_mask] * noise.shape[0])
             noise_mask = noise_mask.to(device)
 
-
     real_model = None
     model_management.load_model_gpu(model)
     real_model = model.model
@@ -82,7 +80,7 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
     model_management.load_controlnet_gpu(control_net_models)
 
     if sampler_name in comfy.samplers.KSampler.SAMPLERS:
-        sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise)
+        sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
     else:
         #other samplers
         pass
@@ -254,7 +252,7 @@ class CheckpointLoaderSimpleSequence:
     def INPUT_TYPES(s):
         return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
                              }}
-    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
+    RETURN_TYPES = ("ORIGINAL_MODEL", "CLIP", "VAE")
     FUNCTION = "load_checkpoint"
 
     CATEGORY = "vid2vid"
diff --git a/sd.py b/sd.py
index 2d18ad6..27762be 100644
--- a/sd.py
+++ b/sd.py
@@ -140,12 +140,12 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
     model = instantiate_from_config(model_config)
     model = load_model_weights(model, sd, verbose=False, load_state_dict_to=load_state_dict_to)
 
-    #with torch.inference_mode(mode=False):
-    model.model.diffusion_model = convert_unet_checkpoint(sd, OmegaConf.create({"model": model_config}))
-    if model_management.xformers_enabled():
-        model.model.diffusion_model.enable_xformers_memory_efficient_attention()
+    with torch.inference_mode(mode=False):
+        model.model.diffusion_model = convert_unet_checkpoint(sd, OmegaConf.create({"model": model_config}))
+        if model_management.xformers_enabled():
+            model.model.diffusion_model.enable_xformers_memory_efficient_attention()
 
-    if fp16:
-        model = model.half()
+    #if fp16:
+    #    model = model.half()
 
     return (ModelPatcher(model), clip, vae, clipvision)
diff --git a/tuneavideo/util.py b/tuneavideo/util.py
index ba8df73..ab39f1d 100644
--- a/tuneavideo/util.py
+++ b/tuneavideo/util.py
@@ -21,7 +21,7 @@ def next_step(model_output: Union[torch.FloatTensor, np.ndarray], timestep: int,
 
 def get_noise_pred_single(latents, t, context, unet):
     latents = rearrange(latents.squeeze(0), "c f h w -> f c h w")
-    noise_pred = unet(latents, t.view(1), context=context)
+    noise_pred = unet(latents, t.view(1).to(latents.device), context=context)
     noise_pred = rearrange(noise_pred.unsqueeze(0), "b f c h w -> b c f h w")
     return noise_pred