import sys import torch import numpy as np import matplotlib.pyplot as plt from pathlib import Path rootdir = Path(sys.path[0]).absolute().parents[1] sys.path.append(str(rootdir)) torch.cuda.set_device(0)