Speed up Dino with DALI about dali HOT 3 OPEN

class Solarize:
    def __init__(self, threshold: int = 128) -> None:
        self._threshold = threshold

    def __call__(self, img):
        inverted_img = 255 - img
        mask = img >= self._threshold
        return mask * inverted_img + (True ^ mask) * img

solarize = Solarize()
            
@pipeline_def(seed=seed, batch_size=1, enable_conditionals=True)
def dino_dali_pipeline(image_dir, local_crops_number=8, device="mixed"):
    
    jpegs, _ = fn.readers.file(file_root=image_dir, random_shuffle=True)

    decoded_jpegs = fn.decoders.image(jpegs, device=device)

    cropped_jpegs = fn.crop(decoded_jpegs, crop=(16384, 16384))
    
    #
    # Global Transform 1
    #
    gt1 = fn.random_resized_crop(cropped_jpegs, size=224, random_aspect_ratio=(1.0, 1.0))

    ## Random Horizontal Flip
    coin = fn.random.coin_flip()
    if coin:
        gt1 = fn.flip(gt1, horizontal=1)

    ## Color Jitter
    coin = fn.random.coin_flip(probability=0.8)
    if coin:
        gt1 = fn.color_twist(gt1, brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)

    ## Random Grayscale
    coin = fn.random.coin_flip(probability=0.2)
    if coin:
        gt1 = fn.color_space_conversion(gt1, image_type=types.RGB, output_type=types.GRAY)

    ## Gaussian Blur
    gt1 = fn.gaussian_blur(gt1, sigma=(0.1, 2.0))

    ## Normalize 
    gt1 = fn.normalize(gt1)

    #
    # Global Transform 2
    #
    gt2 = fn.random_resized_crop(cropped_jpegs, size=224, random_aspect_ratio=(1.0, 1.0))

    ## Random Horizontal Flip
    coin = fn.random.coin_flip()
    if coin:
        gt2 = fn.flip(gt2, horizontal=1)

    ## Color Jitter
    coin = fn.random.coin_flip(probability=0.8)
    if coin:
        gt2 = fn.color_twist(gt2, brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)

    ## Random Grayscale
    coin = fn.random.coin_flip(probability=0.2)
    if coin:
        gt2 = fn.color_space_conversion(gt2, image_type=types.RGB, output_type=types.GRAY)

    ## Gaussian Blur
    coin = fn.random.coin_flip(probability=0.1)
    if coin:
        gt2 = fn.gaussian_blur(gt2, sigma=(0.1, 2.0))

    ## Solarize    
    coin = fn.random.coin_flip(probability=0.1)
    if coin:
        gt2 = fn.cast(solarize(gt2), dtype=types.UINT8)

    gt2 = fn.normalize(gt2)


    #
    # Local Transformations
    #
    
    crops = [gt1, gt2]

    for _ in range(local_crops_number):
        lt = fn.random_resized_crop(cropped_jpegs, size=96, random_aspect_ratio=(0.5, 0.5))

        ## Random Horizontal Flip
        coin = fn.random.coin_flip()
        if coin:
            lt = fn.flip(lt, horizontal=1)
    
        ## Color Jitter
        coin = fn.random.coin_flip(probability=0.8)
        if coin:
            lt = fn.color_twist(lt, brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)
    
        ## Random Grayscale
        coin = fn.random.coin_flip(probability=0.2)
        if coin:
            gt1 = fn.color_space_conversion(lt, image_type=types.RGB, output_type=types.GRAY)
    
        ## Gaussian Blur
        coin = fn.random.coin_flip(probability=0.5)
        if coin:
            lt = fn.gaussian_blur(lt, sigma=(0.1, 2.0))
    
        ## Normalize 
        lt = fn.normalize(lt)
    
    return *crops, 

pipe = dino_dali_pipeline(image_dir, batch_size=4, num_threads=4, device_id=0)
pipe.build()

iterator = DALIGenericIterator(
    pipelines=pipe,
    output_map=["gt1", "gt2", "lt1", "lt2", "lt3", "lt4", "lt5", "lt6", "lt7", "lt8"],
)

for i, (batch,) in enumerate(iterator):
    print(batch)
    break

RuntimeError: [[/opt/dali/dali/pipeline/data/tensor_list.cc:1012](http://localhost:8888/opt/dali/dali/pipeline/data/tensor_list.cc#line=1011)] Assert on "IsDenseTensor()" failed: The batch must be representable as a tensor - it must have uniform shape and be allocated in contiguous memory.
Stacktrace (88 entries):