Implementation in PyTorchΒΆ