Thanks to support I managed to get this working.
/// <summary>
/// Fires before inserting tree node to the tree. Performs the additional actions before the tree node is inserted.
/// </summary>
/// <param name="treeNodeObj">TreeNode object to insert to tree</param>
/// <param name="parentNodeId">Parent node ID.</param>
/// <param name="tree">TreeProvider object to use</param>
public override void OnBeforeInsert(object treeNodeObj, int parentNodeId, object tree)
{
// To change a Document Type's default template dynamically when creating it,
// you will need to use the custom tree node handler (OnBeforeInsert method/event).
// Then change there DocumentPageTemplateID of available TreeNode objects to point to
// some page templage (table CMS_PageTemplate -> Column PageTemplateID).
// In this custom tree node handler you can check if the TreeNode object is of some
// specific class (in this case document type) and if so, then change the page template
// for document. You can find DocumentPageTemplateID in the CMS_Document table.
// PageTemplateID PageTemplateDisplayName
// 25 Product multicolumn view
TreeNode treeNode = (TreeNode)treeNodeObj;
if (treeNode.NodeClassName.ToLower() == "cms.cellphone")
{
treeNode.DocumentPageTemplateID = 25;
}
}