diff --git a/swh/loader/pypi/loader.py b/swh/loader/pypi/loader.py
--- a/swh/loader/pypi/loader.py
+++ b/swh/loader/pypi/loader.py
@@ -36,7 +36,6 @@
 
     def __init__(self, client=None):
         super().__init__(logging_class='swh.loader.pypi.PyPILoader')
-        self.origin_id = None
         if not client:
             temp_directory = self.config['temp_directory']
             os.makedirs(temp_directory, exist_ok=True)
@@ -98,7 +97,7 @@
         self.visit_date = None  # loader core will populate it
 
     def _known_artifacts(self, last_snapshot):
-        """Retrieve the known releases/artifact for the origin_id.
+        """Retrieve the known releases/artifact for the origin.
 
         Args
             snapshot (dict): Last snapshot for the visit
@@ -128,10 +127,10 @@
         """Retrieve the last snapshot
 
         """
-        snapshot = self.storage.snapshot_get_latest(self.origin_id)
-        if snapshot and snapshot.pop('next_branch', None):
-            snapshot = snapshot_get_all_branches(self.storage, snapshot['id'])
-        return snapshot
+        visit = self.storage.origin_visit_get_latest(
+            self.origin, require_snapshot=True)
+        if visit:
+            return snapshot_get_all_branches(self.storage, visit['snapshot'])
 
     def prepare(self, project_name, project_url,
                 project_metadata_url=None):