AOSSIE-Org / AOSSIE-Org/EduAid

[BUG]: PyTorch device mismatch in AnswerPredictor NLI model

Aperta
#442 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug
Lingua principale
JavaScript
Stelle
171
Fork
425
Metriche di merge delle PR
Nessuna PR unita negli ultimi 30g

Descrizione

### 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

Guida per i contributori

Nessuna guida per i contributori indicizzata per questo repository

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.