diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu
index e39c20c..362e0a2 100644
--- a/norch/csrc/cuda.cu
+++ b/norch/csrc/cuda.cu
@@ -40,8 +40,6 @@ __host__ void cuda_to_cpu(Tensor* tensor) {
     const char* device_str = "cpu";
     tensor->device = (char*)malloc(strlen(device_str) + 1);
     strcpy(tensor->device, device_str); 
-
-    printf("Successfully sent tensor to: %s\n", tensor->device);
 }
 
 __global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
diff --git a/train.py b/train.py
index 9ad1187..29d8bc4 100644
--- a/train.py
+++ b/train.py
@@ -1,4 +1,4 @@
-"""import os
+import os
 import norch
 import norch.distributed as dist
 
@@ -30,30 +30,3 @@ def main():
 if __name__ == "__main__":
     main()
 
-"""
-import norch
-import matplotlib.pyplot as plt
-import numpy as np
-import random
-from norch.norchvision import transforms
-
-train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transforms.ToTensor())
-train_sampler = norch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=10, rank=2)
-train_loader = norch.utils.data.Dataloader(train_data, batch_size = 50, sampler=train_sampler)
-input_sample, target_sample = train_data[0]
-
-fig = plt.figure(figsize = (20, 10))
-columns = 4
-rows = 2
-
-# Choose a random image
-for image_index, batch in enumerate(train_loader): 
-
-    image, label = batch
-    fig.add_subplot(rows, columns, image_index+1)
-    plt.imshow(np.array(image).reshape(28, 28))
-    plt.title(label)
-    plt.axis('off')
-    if image_index > 6:
-        break
-plt.show()
\ No newline at end of file