fix for unet hijack breaking the train tab

# Below are monkey patches to enable upcasting a float16 UNet for float32 sampling

def apply_model(orig_func, self, x_noisy, t, cond, **kwargs):

    if isinstance(cond, dict):

        for y in cond.keys():

            cond[y] = [x.to(devices.dtype_unet) if isinstance(x, torch.Tensor) else x for x in cond[y]]

    with devices.autocast():

        return orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs).float()