50 lines
1.2 KiB
Rust
50 lines
1.2 KiB
Rust
//! TDD tests for BERT device handling
|
|
//! These tests define expected behavior for device management
|
|
//!
|
|
//! NOTE: Disabled until BERT API is fully implemented
|
|
|
|
#![cfg(feature = "disabled_tests")]
|
|
|
|
use rtx_tensor::Device;
|
|
use rtx_transformers::architectures::bert_trait_impl::BertModel;
|
|
|
|
#[test]
|
|
fn test_bert_model_device_reference() {
|
|
// Test that BertModel can return a device reference
|
|
let model = BertModel::default();
|
|
|
|
// Should be able to get device reference
|
|
let device = model.device();
|
|
|
|
// Device should be valid
|
|
match device {
|
|
Device::Cuda(_) => {
|
|
// CUDA device is valid
|
|
}
|
|
_ => {
|
|
// Other devices are also valid
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_architecture_type() {
|
|
// Test architecture type identification
|
|
let model = BertModel::default();
|
|
|
|
let arch_type = model.architecture_type();
|
|
assert_eq!(arch_type, "BERT");
|
|
}
|
|
|
|
#[test]
|
|
fn test_bert_model_has_device() {
|
|
// Test that model has a consistent device
|
|
let model = BertModel::default();
|
|
|
|
let device1 = model.device();
|
|
let device2 = model.device();
|
|
|
|
// Should return the same device reference
|
|
assert_eq!(device1, device2);
|
|
}
|