Pytorch tensor.unfold


I had to use tensorrt but the tensor unfold operation was not supported and the output of the converted model was rubbish.
So , I wrote the unfold with simple operations (forum thread) but it is extremely slow.
Do you think python’s for loops is causing this?

Yes, a nested loop could create a large overhead compared to a single operation.

OK thank you , I suppose it is possible to replace unfold with conv2d then and retrain the model then?

Yes, if you were using unfold and a matmul to execute a convolution, you could directly use a conv layer instead.