pkgs/development/python-modules/tensorboardx/default.nix

   1 { boto3, buildPythonPackage, crc32c, fetchFromGitHub, lib, matplotlib, moto
   2 , numpy, pillow, pytorch, protobuf, six, pytestCheckHook
   3 , tensorflow-tensorboard, torchvision }:
   4
   5 buildPythonPackage rec {
   6   pname = "tensorboardx";
   7   version = "2.0";
   8
   9   src = fetchFromGitHub {
  10     owner = "lanpa";
  11     repo = "tensorboardX";
  12     rev = "v${version}";
  13     sha256 = "0qqalq0fhbx0wnd8wdwhyhkkv2brvj9qbk3373vk3wjxbribf5c7";
  14   };
  15
  16   checkInputs = [
  17     pytestCheckHook boto3 crc32c matplotlib moto pillow pytorch tensorflow-tensorboard torchvision
  18   ];
  19
  20   propagatedBuildInputs = [ numpy protobuf six ];
  21
  22   # apparently torch API changed a bit at 1.6
  23   postPatch = ''
  24     substituteInPlace tensorboardX/pytorch_graph.py --replace "torch.onnx.set_training(model, False)" "torch.onnx.select_model_mode_for_export(model, torch.onnx.TrainingMode.EVAL)"
  25   '';
  26
  27
  28   disabledTests = [ "test_TorchVis"  "test_onnx_graph" ];
  29
  30   meta = with lib; {
  31     description = "Library for writing tensorboard-compatible logs";
  32     homepage = "https://github.com/lanpa/tensorboardX";
  33     license = licenses.mit;
  34     maintainers = with maintainers; [ lebastr akamaus ];
  35     platforms = platforms.all;
  36   };
  37 }