AOSSIE-Org / AOSSIE-Org/EduAid

[BUG]: PyTorch device mismatch in AnswerPredictor NLI model

Ouverte
#442 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub
bug
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.

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.