diff --git a/package.json b/package.json
index d7e3227..f67581d 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "react-gojs",
-    "version": "4.5.0",
+    "version": "4.6.0",
     "description": "GoJS React integration",
     "main": "dist/index.js",
     "types": "dist/index.d.ts",
diff --git a/src/GojsDiagram.test.tsx b/src/GojsDiagram.test.tsx
index 4fb6f93..9bcfdfb 100644
--- a/src/GojsDiagram.test.tsx
+++ b/src/GojsDiagram.test.tsx
@@ -7,7 +7,7 @@ import { ModelChangeEventType } from './modelChangeEvent';
 
 const groupName = 'myGroup';
 const singleNode = 'singleNode';
-
+jest.useFakeTimers();
 describe('<GojsDiagram />', () => {
     const portFrom = 'R';
     const portTo = 'L';
@@ -96,6 +96,8 @@ describe('<GojsDiagram />', () => {
     let keyIndex = 0;
 
     beforeEach(() => {
+        Object.defineProperty(Element.prototype, 'clientWidth', { value: 100 });
+        Object.defineProperty(Element.prototype, 'clientHeight', { value: 100 });
         keyIndex = 0;
         const dom = document.body;
         modelChangeCallback = jest.fn();
@@ -133,6 +135,7 @@ describe('<GojsDiagram />', () => {
             />,
             { attachTo: dom }
         );
+        jest.runAllTimers();
     });
 
     it('should default to "category" for nodeCategoryProperty', () => {
@@ -223,9 +226,7 @@ describe('<GojsDiagram />', () => {
 
     testCases.forEach(test => {
         // tslint:disable-next-line:max-line-length
-        it(`should update the render of the diagram (nodes and links) based on the new model provided as prop - case: ${
-            test.description
-        }`, () => {
+        it(`should update the render of the diagram (nodes and links) based on the new model provided as prop - case: ${test.description}`, () => {
             checkIfDiagramRendersModel(model, diagram);
             wrapper.setProps({ model: test.updatedModel });
             checkIfDiagramRendersModel(test.updatedModel, diagram);
diff --git a/src/GojsDiagram.tsx b/src/GojsDiagram.tsx
index 3db8551..237828c 100644
--- a/src/GojsDiagram.tsx
+++ b/src/GojsDiagram.tsx
@@ -42,6 +42,10 @@ export interface GojsModel extends go.Model {
 class GojsDiagram<N extends BaseNodeModel, L extends LinkModel> extends React.PureComponent<GojsDiagramProps<N, L>>
     implements DiagramNotificationDelegate<N, L> {
     private myDiagram: Diagram;
+    private divRef: React.RefObject<HTMLDivElement>;
+
+    private mountInterval;
+
     private eventsToDispatch: ModelChangeEvent<N, L>[];
     private modelChangedHandlers = [
         new AddNodeModelChangedHandler<N, L>(),
@@ -55,19 +59,36 @@ class GojsDiagram<N extends BaseNodeModel, L extends LinkModel> extends React.Pu
 
     constructor(props: GojsDiagramProps<N, L>) {
         super(props);
+        this.divRef = React.createRef();
         this.eventsToDispatch = [];
         this.modelChangedHandler = this.modelChangedHandler.bind(this);
     }
 
     componentDidMount() {
-        this.init();
+        let intervalCount = 0;
+        this.mountInterval = setInterval(() => {
+            if (this.divRef.current && this.divRef.current.clientWidth && this.divRef.current.clientHeight) {
+                this.init();
+                clearInterval(this.mountInterval);
+            } else {
+                if (intervalCount > 10) {
+                    clearInterval(this.mountInterval);
+                }
+                intervalCount++;
+            }
+        }, 10);
     }
 
     componentWillUnmount() {
-        if (this.props.onModelChange) {
+        if (this.myDiagram && this.props.onModelChange) {
             this.myDiagram.removeModelChangedListener(this.modelChangedHandler);
         }
-        this.myDiagram.clear();
+        if (this.myDiagram) {
+            this.myDiagram.clear();
+        }
+        if (this.mountInterval) {
+            clearInterval(this.mountInterval);
+        }
     }
 
     componentDidUpdate() {
@@ -109,7 +130,7 @@ class GojsDiagram<N extends BaseNodeModel, L extends LinkModel> extends React.Pu
         }
     }
     render() {
-        return <div id={this.props.diagramId} className={this.props.className} />;
+        return <div ref={this.divRef} id={this.props.diagramId} className={this.props.className} />;
     }
 
     enqueueEvent(event: ModelChangeEvent<N, L>) {