macOS Sequoia 15.1.1 with MPS and PyTorch (BFloat16 Unsupported)

[SkyForceCoder]

please check this issue

comfyanonymous/ComfyUI#5829

I think the issue is due to the model weight being in bfloat16, and even after converting it, it still detects the weight as bfloat16.
https://github.com/city96/ComfyUI-GGUF/tree/main/tools

I used all the models i see bfloat16 error,
https://huggingface.co/city96/FLUX.1-schnell-gguf/tree/main

is it possible to convert for intel mac AMD

This is not working

% PYTORCH_ENABLE_MPS_FALLBACK=1 PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 python main.py --use-split-cross-attention --force-fp16

Processor:

Model: Intel(R) Xeon(R) W-2140B CPU
Speed: 3.20 GHz
Cores: 8 cores (with Hyper-Threading enabled)

Memory (RAM):

Total RAM: 32 GB (34359738368 bytes)

GPU:

Model: Radeon Pro Vega 56
VRAM: 8 GB
Bus: PCIe (x16)
Metal Support: Metal 3 (supports Apple's graphics framework)

[city96]

Hmmm, the actual weights for the model you linked are all in different quantized formats (the old flux ones use FP16 for the mixed weights because they're from before I had added BF16 support).

I believe FP16 might cause NaN issues with flux, which is why the forward pass is by default cast to bf16 in comfy (and it sort of looks like the fp16 flag isn't respected for flux for some reason).

Could you see if --force-fp32 works as a startup flag?

[srcrs]

me too

[SkyForceCoder]

@city96

% python main.py --force-fp32
[START] Security scan
[DONE] Security scan

ComfyUI-Manager: installing dependencies done.

** ComfyUI startup time: 2024-12-11 21:30:51.125370
** Platform: Darwin
** Python version: 3.10.12 (main, Jul 5 2023, 15:34:07) [Clang 14.0.6 ]
** Python executable: /Volumes/Spidy/VirtualEnvDir/comfyui/bin/python
** ComfyUI Path: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI
** Log path: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfyui.log

Prestartup times for custom nodes:
0.0 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/rgthree-comfy
6.3 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Manager

Total VRAM 32768 MB, total RAM 32768 MB
pytorch version: 2.2.2
Forcing FP32, if this improves things please report it.
Set vram state to: SHARED
Device: mps
Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --use-split-cross-attention
[Prompt Server] web root: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/web
[Crystools INFO] Crystools version: 1.21.0
[Crystools INFO] CPU: Intel(R) Xeon(R) W-2140B CPU @ 3.20GHz - Arch: x86_64 - OS: Darwin 24.1.0
[Crystools ERROR] Could not init pynvml (Nvidia).NVML Shared Library Not Found
[Crystools WARNING] No GPU with CUDA detected.

Loading: ComfyUI-Impact-Pack (V7.14)

Loading: ComfyUI-Impact-Pack (Subpack: V0.8)

[Impact Pack] Wildcards loading done.

Loading: ComfyUI-Manager (V2.55)

ComfyUI Version: v0.3.6-17-g8af9a91 | Released on '2024-12-06'

[ComfyUI-Manager] default cache updated: https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/model-list.json
[ComfyUI-Manager] default cache updated: https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/alter-list.json
[ComfyUI-Manager] default cache updated: https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/custom-node-list.json
[ComfyUI-Manager] default cache updated: https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/extension-node-map.json
[ComfyUI-Manager] default cache updated: https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/github-stats.json
/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/timm/models/layers/init.py:48: FutureWarning: Importing from timm.models.layers is deprecated, please import via timm.layers
warnings.warn(f"Importing from {name} is deprecated, please import via timm.layers", FutureWarning)
Please 'pip install xformers'
Nvidia APEX normalization not installed, using PyTorch LayerNorm
[rvtools INFO] RvTools Version: 2.1.0
[rvtools INFO] RvTools v2 Version: 2.3.5

Comfyroll Studio v1.76 : 175 Nodes Loaded

** For changes, please see patch notes at https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/blob/main/Patch_Notes.md
** For help, please see the wiki at https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki

[rgthree-comfy] Loaded 42 exciting nodes. 🎉

Import times for custom nodes:
0.0 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/websocket_image_save.py
0.0 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI_JPS-Nodes
0.0 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale
0.0 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Custom-Scripts
0.0 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/Derfuu_ComfyUI_ModdedNodes
0.0 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-GGUF
0.0 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/rgthree-comfy
0.0 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI_Comfyroll_CustomNodes
0.1 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/x-flux-comfyui
0.1 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-RvTools_v2
0.1 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-RvTools
0.1 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Long-CLIP
0.4 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Manager
0.5 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/comfyui-tensorops
0.9 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Florence2
1.2 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-PuLID-Flux-Enhanced
2.3 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Crystools
4.9 seconds: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Impact-Pack

Starting server

To see the GUI go to: http://127.0.0.1:8188
FETCH DATA from: /Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Manager/extension-node-map.json [DONE]
[ERROR] An error occurred while retrieving information for the 'Florence2ModelLoader' node.
Traceback (most recent call last):
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/server.py", line 563, in get_object_info
out[x] = node_info(x)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/server.py", line 530, in node_info
info['input'] = obj_class.INPUT_TYPES()
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Florence2/nodes.py", line 150, in INPUT_TYPES
"model": ([item.name for item in Path(folder_paths.models_dir, "LLM").iterdir() if item.is_dir()], {"tooltip": "models are expected to be in Comfyui/models/LLM folder"}),
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Florence2/nodes.py", line 150, in
"model": ([item.name for item in Path(folder_paths.models_dir, "LLM").iterdir() if item.is_dir()], {"tooltip": "models are expected to be in Comfyui/models/LLM folder"}),
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/pathlib.py", line 1017, in iterdir
for name in self._accessor.listdir(self):
FileNotFoundError: [Errno 2] No such file or directory: '/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/models/LLM'

[]
[]
got prompt
Using split attention in VAE
Using split attention in VAE
Requested to load FluxClipModel_
loaded completely 9.5367431640625e+25 4777.53759765625 True
/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-GGUF/nodes.py:79: UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/torch/csrc/utils/tensor_numpy.cpp:212.)
torch_tensor = torch.from_numpy(tensor.data) # mmap

ggml_sd_loader:
GGMLQuantizationType.F16 476
GGMLQuantizationType.Q4_0 304
model weight dtype torch.float32, manual cast: None
model_type FLUX
Requested to load FluxClipModel_
loaded completely 9.5367431640625e+25 4777.53759765625 True
Requested to load Flux
loaded completely 9.5367431640625e+25 6476.5469970703125 True
0%| | 0/20 [00:01<?, ?it/s]
!!! Exception during processing !!! The operator 'aten::rshift.Tensor' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on pytorch/pytorch#77764. As a temporary fix, you can set the environment variable PYTORCH_ENABLE_MPS_FALLBACK=1 to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
Traceback (most recent call last):
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/execution.py", line 324, in execute
output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/execution.py", line 199, in get_output_data
return_values = _map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/execution.py", line 170, in _map_node_over_list
process_inputs(input_dict, i)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/execution.py", line 159, in process_inputs
results.append(getattr(obj, func)(**inputs))
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/nodes.py", line 1467, in sample
return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/nodes.py", line 1434, in common_ksampler
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Impact-Pack/modules/impact/sample_error_enhancer.py", line 22, in informative_sample
raise e
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-Impact-Pack/modules/impact/sample_error_enhancer.py", line 9, in informative_sample
return original_sample(*args, **kwargs) # This code helps interpret error messages that occur within exceptions but does not have any impact on other operations.
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/sample.py", line 43, in sample
samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 1020, in sample
return sample(self.model, noise, positive, negative, cfg, self.device, sampler, sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 918, in sample
return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 904, in sample
output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/patcher_extension.py", line 110, in execute
return self.original(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 873, in outer_sample
output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 857, in inner_sample
samples = executor.execute(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/patcher_extension.py", line 110, in execute
return self.original(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 714, in sample
samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/k_diffusion/sampling.py", line 155, in sample_euler
denoised = model(x, sigma_hat * s_in, **extra_args)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 384, in call
out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 839, in call
return self.predict_noise(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 842, in predict_noise
return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 364, in sampling_function
out = calc_cond_batch(model, conds, x, timestep, model_options)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 200, in calc_cond_batch
return executor.execute(model, conds, x_in, timestep, model_options)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/patcher_extension.py", line 110, in execute
return self.original(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/samplers.py", line 313, in calc_cond_batch
output = model.apply_model(input_x, timestep, **c).chunk(batch_chunks)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/model_base.py", line 128, in apply_model
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/patcher_extension.py", line 110, in execute
return self.original(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/model_base.py", line 157, in _apply_model
model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds).float()
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/ldm/flux/model.py", line 184, in forward
out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control, transformer_options)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/ldm/flux/model.py", line 134, in forward_orig
img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/ldm/flux/layers.py", line 146, in forward
img_mod1, img_mod2 = self.img_mod(vec)
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/ldm/flux/layers.py", line 108, in forward
out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/comfy/ops.py", line 68, in forward
return self.forward_comfy_cast_weights(*args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-GGUF/ops.py", line 152, in forward_comfy_cast_weights
return self.forward_ggml_cast_weights(input, *args, **kwargs)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-GGUF/ops.py", line 174, in forward_ggml_cast_weights
weight, bias = self.cast_bias_weight(input)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-GGUF/ops.py", line 146, in cast_bias_weight
weight = s.get_weight(s.weight.to(device), dtype)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-GGUF/ops.py", line 119, in get_weight
weight = dequantize_tensor(tensor, dtype, self.dequant_dtype)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-GGUF/dequant.py", line 23, in dequantize_tensor
return dequantize(tensor.data, qtype, oshape, dtype=dequant_dtype).to(dtype)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-GGUF/dequant.py", line 43, in dequantize
blocks = dequantize_blocks(blocks, block_size, type_size, dtype)
File "/Volumes/Spidy/GitHub/AI/ComfyUIRoot/ComfyUI/custom_nodes/ComfyUI-GGUF/dequant.py", line 117, in dequantize_blocks_Q4_0
qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1))
File "/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/_tensor.py", line 1418, in torch_function
ret = func(*args, **kwargs)
NotImplementedError: The operator 'aten::rshift.Tensor' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on pytorch/pytorch#77764. As a temporary fix, you can set the environment variable PYTORCH_ENABLE_MPS_FALLBACK=1 to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.

[SkyForceCoder]

@city96

if used PYTORCH_ENABLE_MPS_FALLBACK=1 then it work on CPU but very slow

% PYTORCH_ENABLE_MPS_FALLBACK=1 python main.py --force-fp32

Requested to load FluxClipModel_
loaded completely 9.5367431640625e+25 4777.53759765625 True
Requested to load Flux
loaded completely 9.5367431640625e+25 6476.5469970703125 True
0%| | 0/20 [00:00<?, ?it/s]/Volumes/Spidy/VirtualEnvDir/comfyui/lib/python3.10/site-packages/torch/_tensor.py:1418: UserWarning: The operator 'aten::rshift.Tensor' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:13.)
ret = func(*args, **kwargs)
50%|████████████████████████████████████████████████████████████████ | 10/20 [02:35<02:35, 15.57s/it]

[city96]

That error isn't about BFloat16 though, it's about rshift, which is a completely different issue. Based on this comment the rshift thing should work on latest pytorch: #27 (comment)

You're using pytorch 2.2.2 which is probably why it's failing.