Update core/data_architecture/tensor_core_layers.py

core/data_architecture/tensor_core_layers.py

@@ -2,6 +2,21 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.cuda.amp import autocast
+from torch.optim import Adam
+import cupy as cp  
+import cudf 
+import flash_attn  
+import onnx
+import onnxruntime as ort
+import tensorrt as trt
+from nemo.collections.nlp.models import GPTModel
+from nemo.collections.tts.models import FastPitchModel
+from nemo.collections.asr.models import EncDecCTCModel
+from torch2trt import torch2trt  # Convert PyTorch to TensorRT
+from transformers import AutoModel, AutoTokenizer
+import apex
+from apex import amp
+from apex.optimizers import FusedAdam
 
 class TensorCoreLinear(nn.Module):
     """