eora-ai / torchok

Production-oriented Computer Vision models training pipeline for common tasks: classification, segmentation, detection and representationšŸ„¤
https://torchok.readthedocs.org
Apache License 2.0
54 stars 6 forks source link

Representation Metrics in DDP mode #149

Open PososikTeam opened 2 years ago

PososikTeam commented 2 years ago

Describe the bug In DDP mode representation metrics doesn't work. Due to in update function gpu embedding convert to cpu embedding and after that in compute method in DDP mode called all_gather() which work only with gpu tensors.