From 9b9cf9469ffad6c98efa980c4e12ea6d3fee4e06 Mon Sep 17 00:00:00 2001
From: Dmitry Vyal <dmitryvyal@gmail.com>
Date: Sat, 5 Sep 2020 23:46:04 +0300
Subject: [PATCH] python3Packages.tensorboardx: fix build after upgrade to
 PyTorch 1.6

---
 pkgs/development/python-modules/tensorboardx/default.nix | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pkgs/development/python-modules/tensorboardx/default.nix b/pkgs/development/python-modules/tensorboardx/default.nix
index d2532777d00..a95f8e02928 100644
--- a/pkgs/development/python-modules/tensorboardx/default.nix
+++ b/pkgs/development/python-modules/tensorboardx/default.nix
@@ -19,6 +19,12 @@ buildPythonPackage rec {
 
   propagatedBuildInputs = [ numpy protobuf six ];
 
+  # apparently torch API changed a bit at 1.6
+  postPatch = ''
+    substituteInPlace tensorboardX/pytorch_graph.py --replace "torch.onnx.set_training(model, False)" "torch.onnx.select_model_mode_for_export(model, torch.onnx.TrainingMode.EVAL)"
+  '';
+
+
   disabledTests = [ "test_TorchVis"  "test_onnx_graph" ];
 
   meta = with lib; {