Thanks, I’ll try to file a bug.
I’m having similar issues with attribution, I am not able to uncomment below without running out of memory, even with gc:
    # following the tutorial on attribution for semantic segmentation
    def save_attribution(label, target):
        lc_attr = layer_cond.attribute(input, target=target, n_steps=5, internal_batch_size=1)
        fig, ax = viz.visualize_image_attr_multiple(
            lc_attr[0].cpu().permute(1,2,0).detach().numpy(),
            original_image=orig,
            signs=["positive", "negative"],
            methods=["blended_heat_map", "blended_heat_map"],
            use_pyplot=False
        )
        plt.savefig(f'plot_{label}.png', bbox_inches='tight')
        del lc_attr, fig, ax
        torch.cuda.empty_cache()
        gc.collect()


    save_attribution('smooth', 1)
    # save_attribution('grass', 2)
    # save_attribution('rough', 3)