fix: only use .cpu() in the final step in collector.py to reduce the usage of cpu#2071