AOSSIE-Org / AOSSIE-Org/EduAid

[BUG]: PyTorch device mismatch in AnswerPredictor NLI model

Abierto
#442 0 comentarios 0 reacciones 0 asignados Ver en GitHub
bug
Lenguaje dominante
JavaScript
Estrellas
171
Forks
425
Métricas de merge de PR
Sin PR fusionados en 30 d

Descripción

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

Guía de contribución

No hay ninguna guía de contribución indexada para este repositorio

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.