Updated README

README.md

@@ -49,13 +49,13 @@ audio, sr = torchaudio.load("path/to/audio.wav")
 # attention_mask[0, 8000:] = 0  # Example: mask second half of first sample
 
 # Method 1: End-to-end processing (encode + decode)
-with torch.no_grad():
+with torch.no_grad(), torch.autocast(device_type='cuda')::
     outputs = model(audio)  # Optionally pass attention_mask=attention_mask
     reconstructed_audio = outputs["audio"]
     embeddings = outputs['embeddings']
 
 # Method 2: Separate encoding and decoding
-with torch.no_grad():
+with torch.no_grad(), torch.autocast(device_type='cuda')::
     # Encode audio to embeddings
     embeddings = model.encode(audio)  # Optionally pass attention_mask=attention_mask