fix versioning, bump version

Author: 1adrianb

README.md

@@ -1,4 +1,4 @@
-[![License](https://img.shields.io/badge/License-BSD%203--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)  [![Build Status](https://travis-ci.com/1adrianb/pytorch-estimate-flops.svg?branch=master)](https://travis-ci.com/1adrianb/pytorch-estimate-flops)
+[![License](https://img.shields.io/badge/License-BSD%203--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)  [![Test Pytorch Flops Counter](https://github.com/1adrianb/pytorch-estimate-flops/workflows/Test%20Pytorch%20Flops%20Counter/badge.svg)](https://travis-ci.com/1adrianb/pytorch-estimate-flops)
 [![PyPI](https://img.shields.io/pypi/v/pthflops.svg?style=flat)](https://pypi.org/project/pthflops/)
 
 # pytorch-estimate-flops

pthflops/__init__.py

@@ -1,3 +1,3 @@
 from .ops import count_ops
 
-__version__ = '0.3.4'
+__version__ = '0.3.5'

pthflops/ops.py

@@ -309,7 +309,7 @@ def count_ops(model, input, custom_ops={}, ignore_layers=[], print_readable=True
         if LooseVersion(torch.__version__) >= LooseVersion('1.6.0') and \
                 LooseVersion(torch.__version__) < LooseVersion('1.8.0'):
             version = 2
-        else:
+        elif LooseVersion(torch.__version__) >= LooseVersion('1.8.0'):
             version = 3
     else:
         # PyTorch 1.3 and bellow