diff --git a/tests/test_vad.py b/tests/test_vad.py
index 5ff5fa1..cfc2d3a 100644
--- a/tests/test_vad.py
+++ b/tests/test_vad.py
@@ -1,13 +1,12 @@
 import unittest
 import numpy as np
-import torch
 from whisper_live.tensorrt_utils import load_audio
-from whisper_live.vad import VoiceActivityDetection
+from whisper_live.vad import VoiceActivityDetector
 
 
 class TestVoiceActivityDetection(unittest.TestCase):
     def setUp(self):
-        self.vad = VoiceActivityDetection()
+        self.vad = VoiceActivityDetector()
         self.sample_rate = 16000
 
     def generate_silence(self, duration_seconds):
@@ -18,10 +17,10 @@ class TestVoiceActivityDetection(unittest.TestCase):
 
     def test_vad_silence_detection(self):
         silence = self.generate_silence(3)
-        speech_prob = self.vad(torch.from_numpy(silence.copy()), self.sample_rate).item()
-        self.assertLess(speech_prob, 0.5, "VAD incorrectly identified silence as speech.")
+        is_speech_present = self.vad(silence.copy())
+        self.assertFalse(is_speech_present, "VAD incorrectly identified silence as speech.")
 
     def test_vad_speech_detection(self):
-        audio_tensor = torch.from_numpy(load_audio("assets/jfk.flac"))
-        speech_prob = self.vad(audio_tensor, self.sample_rate).item()
-        self.assertGreater(speech_prob, 0.5, "VAD failed to identify speech segment.")
+        audio_tensor = load_audio("assets/jfk.flac")
+        is_speech_present = self.vad(audio_tensor)
+        self.assertTrue(is_speech_present, "VAD failed to identify speech segment.")