import torchshard as ts