TypeError: Got unsupported ScalarType BFloat16

[yannick3751]

Using Google Colab
This issue occurs with the A100 GPU, L4 GPU, T4 GPU and TPU v2-8
Everything works as normal (however slower) with the regular CPU runtime

###########################################

sam.generate(image, output="masks.tif", foreground=True, unique=True)
delivers the TypeError: Got unsupported ScalarType BFloat16

###########################################

TypeError Traceback (most recent call last)

in <cell line: 34>()
32
33 # Perform segmentation
---> 34 GeoSAM(input_path, output_path)

5 frames

in GeoSAM(image_path, output_path)
16
17 if sam_version == 1:
---> 18 sam.generate(image, output="masks.tif", foreground=True, unique=True)
19 sam.tiff_to_vector("masks.tif", output_path)
20

/usr/local/lib/python3.10/dist-packages/samgeo/samgeo.py in generate(self, source, output, foreground, batch, batch_sample_size, batch_nodata_threshold, nodata_value, erosion_kernel, mask_multiplier, unique, **kwargs)
219 self.image = image # Store the input image as a numpy array
220 mask_generator = self.mask_generator # The automatic mask generator
--> 221 masks = mask_generator.generate(image) # Segment the input image
222 self.masks = masks # Store the masks as a list of dictionaries
223 self.batch = False

/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py in decorate_context(*args, **kwargs)
114 def decorate_context(*args, **kwargs):
115 with ctx_factory():
--> 116 return func(*args, **kwargs)
117
118 return decorate_context

/usr/local/lib/python3.10/dist-packages/segment_anything/automatic_mask_generator.py in generate(self, image)
161
162 # Generate masks
--> 163 mask_data = self._generate_masks(image)
164
165 # Filter small disconnected regions and holes in masks

/usr/local/lib/python3.10/dist-packages/segment_anything/automatic_mask_generator.py in _generate_masks(self, image)
220 data.filter(keep_by_nms)
221
--> 222 data.to_numpy()
223 return data
224

/usr/local/lib/python3.10/dist-packages/segment_anything/utils/amg.py in to_numpy(self)
73 for k, v in self._stats.items():
74 if isinstance(v, torch.Tensor):
---> 75 self._stats[k] = v.detach().cpu().numpy()
76
77

TypeError: Got unsupported ScalarType BFloat16

[SayiKrishnaA]

.detach().cpu().to(torch.float32).numpy()

Doing this solved the issue in colab