Method set_device

funasr/models/transducer/rnnt_decoder.py:192–199

Set GPU device to use. Args: device: Device ID.

(self, device: torch.device)

        return dec_out.squeeze(1), states

    def set_device(self, device: torch.device) -> None:
        """Set GPU device to use.

        Args:
            device: Device ID.

        """
        self.device = device

    def init_state(self, batch_size: int) -> Tuple[torch.Tensor, Optional[torch.tensor]]:
        """Initialize decoder states.