diff --git a/ebay_api.py b/ebay_api.py
index 6c972d8..223048b 100644
--- a/ebay_api.py
+++ b/ebay_api.py
@@ -154,13 +154,13 @@ class CurateData:
 
     def to_training(self, data):
         raw_df = self.data_frame(data)
-        interm_df1 = raw_df[['ItemID', 'PictureURL', 'PrimaryCategoryID', 'PrimaryCategoryName', 'Title', 'ItemSpecifics']]
-        interm_df1[['ItemID', 'PrimaryCAegoryID']] = interm_df1[['ItemID', 'PrimaryCategoryID']].astype(str)
+        interm_df1 = raw_df.loc[:,['ItemID', 'PictureURL', 'PrimaryCategoryID', 'PrimaryCategoryName', 'Title', 'ItemSpecifics']]
+        interm_df1[['ItemID', 'PrimaryCAegoryID']] = interm_df1.loc[:, ['ItemID', 'PrimaryCategoryID']].astype(str)
         training = interm_df1
         return training
 
     def class_training(self, training):
-        class_df = training[['PictureURL', 'PrimaryCategoryID']]
+        class_df = training.loc[:, ['PictureURL', 'PrimaryCategoryID']]
         return class_training
 
     def nvl_training(self, training):