You can refer to tao_pytorch_backend/nvidia_tao_pytorch/cv/pose_classification/model/st_gcn.py at 9c2d94c0635b1117edfea85a94a6e3d0ead53754 · NVIDIA/tao_pytorch_backend · GitHub and modify it to meet your key points.
Similar topic is PoseClassificationNet Model training on custom dataset.