fix: audio_length_total only accumulated for CUDA kernel

[andrewwhitecdw]

Problem

fix: audio_length_total only accumulated for CUDA kernel

Fix

Replace:

        vram_used_original_total += torch.cuda.max_memory_allocated(device="cuda")
        
        del data, audio_original
        torch.cuda.empty_cache()

with:

        vram_used_original_total += torch.cuda.max_memory_allocated(device="cuda")
        audio_length_total += audio_original.shape[-1]
        
        del data, audio_original
        torch.cuda.empty_cache()

Files changed

• tests/test_cuda_vs_torch_model.py