def METHOD_NAME(self, timestep, label_image, output_dir, filename_zero_padding=3):
    """
    Save a single frame to a 2D or 3D tif
    """
    filename = os.path.join(output_dir, f"mask{timestep:0{filename_zero_padding}}.tif")
    # import ipdb; ipdb.set_trace()
    label_image = np.swapaxes(label_image.squeeze(), 0, 1)
    if len(label_image.shape) == 2:  # 2d
        vigra.impex.writeImage(label_image.astype("uint16"), filename)
    elif len(label_image.shape) == 3:  # 3D
        label_image = np.transpose(label_image, axes=[2, 0, 1])
        tifffile.imwrite(filename, label_image.astype("uint16"))
    else:
        raise ValueError("Image had the wrong dimensions, can only save 2 or 3D images with a single channel")