NVIDIA / TensorRT-Incubator

Experimental projects related to TensorRT
81 stars 12 forks source link

[Tripy] Mixed module list/dict for `tp.Module` #365

Closed farazkh80 closed 1 week ago

farazkh80 commented 1 week ago

PR for addressing issue https://github.com/NVIDIA/TensorRT-Incubator/issues/302. Supporting mixed types is important, since tripy does not have modules like Identity or AveragePool. With this, users can make sequential nets or module lists that include lambda functions or other callables like tp.avgpool.

farazkh80 commented 1 week ago

Thanks for the review @pranavm-nvidia , addressed all PR changes. Should be good for a second pass