MahmoudAshraf97

import numpy as np
import torch
import tensorrt as trt
from polygraphy.backend.trt import CreateConfig, Profile


def build_encoder_t_gather_engine(encoder_output_shape):
    """
    Build a TensorRT engine that:
        input_0: encoder_output, shape (B, T, H), dtype float32

MahmoudAshraf97

from typing import Tuple

import line_profiler
import torch
import torch.nn as nn


class Encoder(nn.Module):
    def __init__(self, feature_size, filter_length, hop_length):
        super(Encoder, self).__init__()

MahmoudAshraf97

import os
import timeit

print(f"Cores: {os.sched_getaffinity(0)}")
import cupy as cp
import numpy as np
import torch


def stft(

MahmoudAshraf97

import torch
import torch.nn as nn
import torch.nn.functional as F


class STFT(nn.Module):
    def __init__(self, filter_length, hop_length):
        super(STFT, self).__init__()
        self.filter_length = filter_length
        self.hop_length = hop_length

MahmoudAshraf97

import torch
import torch.nn as nn
import torch.nn.functional as F


class STFT(nn.Module):
    def __init__(self, filter_length, hop_length):
        super(STFT, self).__init__()
        self.filter_length = filter_length
        self.hop_length = hop_length

MahmoudAshraf97

import os
import torch
from torchaudio.models import wav2vec2_model
from transformers import Wav2Vec2ForCTC, Wav2Vec2CTCTokenizer, AutoConfig
import json
import argparse
import tempfile

# Initialize parser
parser = argparse.ArgumentParser()

MahmoudAshraf97

import torch
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
import numpy as np
import cv2
from datetime import datetime
import skvideo.io

import json