diff --git a/src/modules/partition/DeviceModel.cpp b/src/modules/partition/DeviceModel.cpp
index eafc722de..a401f39a7 100644
--- a/src/modules/partition/DeviceModel.cpp
+++ b/src/modules/partition/DeviceModel.cpp
@@ -21,18 +21,6 @@
 // CalaPM
 #include <core/device.h>
 
-DeviceModel::DeviceInfo::DeviceInfo( Device* dev )
-    : device( dev )
-    , partitionModel( new PartitionModel )
-{
-    partitionModel->init( dev );
-}
-
-DeviceModel::DeviceInfo::~DeviceInfo()
-{
-    delete partitionModel;
-}
-
 DeviceModel::DeviceModel( QObject* parent )
     : QAbstractListModel( parent )
 {
@@ -46,11 +34,7 @@ void
 DeviceModel::init( const QList< Device* >& devices )
 {
     beginResetModel();
-    m_devices.clear();
-    for ( auto device : devices )
-    {
-        m_devices << new DeviceInfo( device );
-    }
+    m_devices = devices;
     endResetModel();
 }
 
@@ -69,7 +53,7 @@ DeviceModel::data( const QModelIndex& index, int role ) const
         return QVariant();
     }
 
-    Device* device = m_devices.at( row )->device;
+    Device* device = m_devices.at( row );
 
     switch ( role )
     {
@@ -87,13 +71,13 @@ DeviceModel::data( const QModelIndex& index, int role ) const
     }
 }
 
-PartitionModel*
-DeviceModel::partitionModelForIndex( const QModelIndex& index ) const
+Device*
+DeviceModel::deviceForIndex( const QModelIndex& index ) const
 {
     int row = index.row();
     if ( row < 0 || row >= m_devices.count() )
     {
         return nullptr;
     }
-    return m_devices.at( row )->partitionModel;
+    return m_devices.at( row );
 }
diff --git a/src/modules/partition/DeviceModel.h b/src/modules/partition/DeviceModel.h
index fb4ced9c2..85b333455 100644
--- a/src/modules/partition/DeviceModel.h
+++ b/src/modules/partition/DeviceModel.h
@@ -40,17 +40,10 @@ public:
     int rowCount( const QModelIndex& parent = QModelIndex() ) const Q_DECL_OVERRIDE;
     QVariant data( const QModelIndex& index, int role = Qt::DisplayRole ) const Q_DECL_OVERRIDE;
 
-    PartitionModel* partitionModelForIndex( const QModelIndex& index ) const;
+    Device* deviceForIndex( const QModelIndex& index ) const;
 
 private:
-    struct DeviceInfo
-    {
-        DeviceInfo( Device* dev );
-        ~DeviceInfo();
-        Device* device;
-        PartitionModel* partitionModel;
-    };
-    QList< DeviceInfo* > m_devices;
+    QList< Device* > m_devices;
 };
 
 #endif /* DEVICEMODEL_H */
diff --git a/src/modules/partition/PartitionCoreModule.cpp b/src/modules/partition/PartitionCoreModule.cpp
index 84b83535b..578c2b122 100644
--- a/src/modules/partition/PartitionCoreModule.cpp
+++ b/src/modules/partition/PartitionCoreModule.cpp
@@ -19,12 +19,28 @@
 #include <PartitionCoreModule.h>
 
 #include <DeviceModel.h>
+#include <PartitionModel.h>
 
 // CalaPM
 #include <CalaPM.h>
 #include <backend/corebackend.h>
 #include <backend/corebackendmanager.h>
 
+
+//- DeviceInfo --------------------------------------------
+PartitionCoreModule::DeviceInfo::DeviceInfo( Device* dev )
+    : device( dev )
+    , partitionModel( new PartitionModel )
+{
+    partitionModel->init( dev );
+}
+
+PartitionCoreModule::DeviceInfo::~DeviceInfo()
+{
+    delete partitionModel;
+}
+
+//- PartitionCoreModule -----------------------------------
 PartitionCoreModule::PartitionCoreModule( QObject* parent )
     : QObject( parent )
     , m_deviceModel( new DeviceModel( this ) )
@@ -34,16 +50,20 @@ PartitionCoreModule::PartitionCoreModule( QObject* parent )
     {
         qFatal( "Failed to init CalaPM" );
     }
+
     CoreBackend* backend = CoreBackendManager::self()->backend();
-    m_devices = backend->scanDevices();
+    QList< Device* > lst = backend->scanDevices();
+    m_deviceModel->init( lst );
+    for ( auto device : lst )
+    {
+        m_devices << new DeviceInfo( device );
+    }
 
-    m_deviceModel->init( m_devices );
 }
 
-QList< Device* >
-PartitionCoreModule::devices() const
+PartitionCoreModule::~PartitionCoreModule()
 {
-    return m_devices;
+    qDeleteAll( m_devices );
 }
 
 DeviceModel*
@@ -51,3 +71,16 @@ PartitionCoreModule::deviceModel() const
 {
     return m_deviceModel;
 }
+
+PartitionModel*
+PartitionCoreModule::partitionModelForDevice( Device* device ) const
+{
+    for ( auto it : m_devices )
+    {
+        if ( it->device == device )
+        {
+            return it->partitionModel;
+        }
+    }
+    return nullptr;
+}
diff --git a/src/modules/partition/PartitionCoreModule.h b/src/modules/partition/PartitionCoreModule.h
index 405210d6d..ae7f6e780 100644
--- a/src/modules/partition/PartitionCoreModule.h
+++ b/src/modules/partition/PartitionCoreModule.h
@@ -24,6 +24,7 @@
 
 class Device;
 class DeviceModel;
+class PartitionModel;
 
 /**
  * Owns the Qt models and the PM devices
@@ -32,13 +33,21 @@ class PartitionCoreModule : public QObject
 {
 public:
     PartitionCoreModule( QObject* parent = nullptr );
-
-    QList< Device* > devices() const;
+    ~PartitionCoreModule();
 
     DeviceModel* deviceModel() const;
 
+    PartitionModel* partitionModelForDevice( Device* device ) const;
+
 private:
-    QList< Device* > m_devices;
+    struct DeviceInfo
+    {
+        DeviceInfo( Device* dev );
+        ~DeviceInfo();
+        Device* device;
+        PartitionModel* partitionModel;
+    };
+    QList< DeviceInfo* > m_devices;
     DeviceModel* m_deviceModel;
 
     void listDevices();
diff --git a/src/modules/partition/PartitionPage.cpp b/src/modules/partition/PartitionPage.cpp
index 67d4cf711..2f3883ae6 100644
--- a/src/modules/partition/PartitionPage.cpp
+++ b/src/modules/partition/PartitionPage.cpp
@@ -39,7 +39,8 @@ PartitionPage::PartitionPage( QWidget* parent )
     connect( m_ui->deviceListView->selectionModel(), &QItemSelectionModel::currentChanged,
              [ this ]( const QModelIndex& index, const QModelIndex& oldIndex )
     {
-        PartitionModel* model = m_core->deviceModel()->partitionModelForIndex( index );
+        Device* device = m_core->deviceModel()->deviceForIndex( index );
+        PartitionModel* model = m_core->partitionModelForDevice( device );
         m_ui->partitionListView->setModel( model );
     } );
 }