Abdelrahman Ahmad

Hello, I am trying to run an onnex model using tensorrt backend, but I get the following error. KeyError: 'output1_before_shuffle' ``` model = onnx.load(args.files) onnx.checker.check_model(model) input_shapes = [[d.dim_value for d...

This update is eliminate the following warning ``` /miniconda/envs/Research/lib/python3.12/site-packages/colt5_attention/topk.py:9: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead. @autocast(enabled = False) ```

Update autocast function from deprecated argument to the new one.