AOSSIE-Org / AOSSIE-Org/EduAid
[BUG]: PyTorch device mismatch in AnswerPredictor NLI model
- Langage dominant
- JavaScript
- Étoiles
- 171
- Forks
- 425
- Métriques de merge des PR
- Aucune PR mergée en 30 j
Description
### Bug Description
### Description
The distilbert-base-uncased-mnli model in the AnswerPredictor class is not explicitly moved to self.device during initialization. Additionally, input tensors in predict_boolean_answer() default to CPU and are not aligned with the model’s device.
### ⚠️ Why This Is a Problem
If the backend is executed on a CUDA-enabled system:
- The model may remain on CPU
- Input tensors may remain on CPU even if future changes move the model to GPU
- This can lead to:
- RuntimeError: Expected all tensors to be on the same device
- Silent CPU fallback (GPU not utilized despite availability)
- Inconsistent hardware behavior across different models in the system
While the current implementation may work in CPU-only environments, it lacks explicit device synchronization and may cause instability in GPU deployments.
### ✅ Expected Behavior
- The NLI model should be explicitly moved to self.device during initialization.
- All input tensors should be moved to the same device before inference.
- Device handling should be consistent with other transformer models in main.py.
### Proposed Fix
- Move the NLI model to self.device after loading.
- Move tokenizer outputs to self.device before inference.
I have implemented and tested this fix locally and will open a PR shortly.
Medium - Feature works but has issues
### Code of Conduct
- [x] I have joined the [Discord server](https://discord.gg/hjUhu33uAn) and will post updates there
- [x] I have searched existing issues to avoid duplicates
Guide de contribution
Aucun guide de contribution indexé pour ce dépôt
Évaluation
Cette issue n'a pas encore été évaluée.