Use torchvision in a different way to avoid type errors

[adamtheturtle]

No description provided.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 56.50%. Comparing base (b47c313) to head (cdeb0b3).
Report is 506 commits behind head on main.

❗ There is a different number of reports uploaded between BASE (b47c313) and HEAD (cdeb0b3). Click for more details.

HEAD has 79 uploads less than BASE

Flag	BASE (b47c313)	HEAD (cdeb0b3)
	80	1

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2140       +/-   ##
===========================================
- Coverage   99.89%   56.50%   -43.40%     
===========================================
  Files          74       74               
  Lines        4999     5001        +2     
  Branches      627      627               
===========================================
- Hits         4994     2826     -2168     
- Misses          5     2106     +2101     
- Partials        0       69       +69     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.